Improve formatting of @overload-decorated function groups in stub files, especially those with docstrings

[AlexWaygood]

⚠️ I'm pretty unfamiliar with Black's codebase, so this patch was mostly authored by Claude. I verified that it has extensive test coverage, that it has the effect that I want, and that it's well documented, but I can't say much about whether it's making the change in the right way or not.

Description

In .pyi stub files, Black currently has no special handling for groups of consecutive decorated functions that share the same name. This comes up quite often with @overload groups, and it results in especially unfortunate formatting if one of the @overload-decorated functions has a docstring in it. Overloads can end up with inconsistent blank lines — sometimes separated from each other, sometimes not — depending on the surrounding context and whether any overloads have docstrings:

# Before (status quo)

@overload                                                                                                                                       
def foo(x: int) -> int:                                         
    """Docs."""

@overload
def foo(x: str) -> str: ...  # a blank line is enforced between the `foo` overloads due to the docstring
def bar(x): ...              # a blank line is removed between the last `foo` overload and `bar`

class MyClass:
    attr: int
    @overload                # no blank line after attr
    def method(self, x: int) -> int: ...
    @overload
    def method(self, x: str) -> str: ...
    attr2: str               # no blank line before attr2

This PR adds a new preview feature (pyi_overload_group_blank_lines) that enforces two rules for decorated function statements in .pyi files:

1. Before a decorated function: always insert a blank line, unless the immediately preceding statement is a decorated function with the same name (i.e. part of the same overload group), or it is the first statement in the block.
2. After a decorated function: always insert a blank line, unless the immediately following statement is a decorated function with the same name.

These rules apply regardless of whether the adjacent statement is a function definition, a class, or a simple statement like a variable annotation.

# After (with `--preview`)

@overload
def foo(x: int) -> int:
    """Docs."""
@overload
def foo(x: str) -> str: ...  # no blank line in between the `foo` overloads

def bar(x): ...              # blank line after `foo` group

class MyClass:
    attr: int

    @overload                # blank line before group
    def method(self, x: int) -> int: ...
    @overload
    def method(self, x: str) -> str: ...

    attr2: str               # blank line after group

The implementation tracks the most recently seen decorated function's name and depth in EmptyLineTracker, and uses this state to determine whether to collapse or separate adjacent statements.

Motivation

The motivation here is that ty codemods docstrings into its vendored copy of typeshed's stubs, and then reformats the stubs with Black after the codemod. We use Black to reformat the stubs because typeshed uses Black upstream, and we want to minimze the diff between our modified version of typeshed and the upstream version of typeshed. Black's formatting also has the nice property that it will never reformat a line with a # type: ignore comment, which means that it's easy to run typeshed's tests on our codemodded version of typeshed and verify that our codemodding didn't break anything in typeshed.

I ran Black from this branch on ty's vendored typeshed stubs, and you can see the result here -- I think it's a big improvement, personally! https://github.com/astral-sh/ruff/compare/main...AlexWaygood:ruff:reformat-stubs?expand=1.

Fixes astral-sh/ty#2945

Checklist - did you ...

• Implement any code style changes under the --preview style, following the stability policy?
• Add an entry in CHANGES.md if necessary?
• Add / update tests if necessary?
• Add new / update outdated documentation?

[github-actions]

diff-shades results comparing this PR (9e23f32) to main (c014b22):

--preview style (View full diff):

╭───────────────────────── Summary ─────────────────────────╮
│ 4 projects & 593 files changed / 3046 changes [+3006/-40] │
│                                                           │
│ ... out of 2 892 016 lines, 13 649 files & 22 projects    │
╰───────────────────────────────────────────────────────────╯
Differences found.

--stable style: no changes

---

What is this? | Workflow run | diff-shades documentation

[AlexWaygood]

diff-shades results comparing this PR (9540cbc) to main (457320a):

This looks about as-expected! Most stubs in the wild don't have docstrings in them, so on most codebases this will result in blank lines being added before or after decorated functions and overload groups. In my opinion, that makes things more readable.

[cobaltt7]

Thanks for this change! I agree that the current heuristic isn't ideal, and this definitely looks like an improvement.
I agree that overloaded functions with the same name should be grouped together, but does it make sense to apply this to all decorators? Otherwise, we end up with code like this - multiple undecorated functions grouped together, then groups of one decorated function each:

It might make more sense to only add the padding if there's multiple functions with the same name, not only if it's decorated. (if this is changed, it should be solely based on the count; I wouldn't want to special-case @overload decorators) But I don't have a strong opinion either way, and the current implementation looks good!
From a skim, the code changes also look good!

[AlexWaygood]

I agree that overloaded functions with the same name should be grouped together, but does it make sense to apply this to all decorators? Otherwise, we end up with code like this - multiple undecorated functions grouped together, then groups of one decorated function each:

Jelle gave me similar feedback offline — since you both gave the same feedback independently, and this was the design aspect I was least sure of, I'll change this. Thanks!

[AlexWaygood]

I'd normally use a set for this kind of containment check, but IIRC mypyc optimizes tuples much better than sets (though maybe that's outdated lore by now...? Not sure...), and I know black compiles with mypyc

[AlexWaygood]

I agree that overloaded functions with the same name should be grouped together, but does it make sense to apply this to all decorators?

Okay, I addressed this point, and also fixed a whole bunch of edge cases thrown up by diff-shades. The heuristics hopefully now mean that the patch only changes formatting in .pyi files for cases involving @overload groups or properties with setters. I also updated the branch of Ruff where I applied this branch's formatting to ty's vendored typeshed stubs: https://github.com/astral-sh/ruff/compare/main...AlexWaygood:ruff:reformat-stubs?expand=1

[AlexWaygood]

Changing the heuristic so that we only add blank lines around decorated functions if there's a sequence of >=2 with the same name has added a fair amount of complexity to this patch. If this is now too complex to stomach, one way to simplify would be to remove some of the handling that I've added to take care of decorated-function groups where one of the decorated functions is inside an if sys.version_info block or similar and others are not.

Possible simplification

diff --git a/src/black/lines.py b/src/black/lines.py
index da54666..96a0ff1 100644
--- a/src/black/lines.py
+++ b/src/black/lines.py
@@ -683,93 +683,6 @@ def _is_start_of_decorated_group(line: Line) -> bool:
             return False
         return EmptyLineTracker._decorated_node_starts_group(decorated_node)
 
-    @staticmethod
-    def _get_block_first_decorated_funcname(line: Line) -> str | None:
-        """Return the function name of the first decorated function in a block.
-
-        *line* must be a block-opening line (ending with ``:``) such as
-        ``if ...:``.  Returns ``None`` when the block doesn't start with a
-        decorated function.
-        """
-        if not line.leaves or line.leaves[-1].type != token.COLON:
-            return None
-        suite = line.leaves[-1].next_sibling
-        if suite is None or not isinstance(suite, Node) or suite.type != syms.suite:
-            return None
-        for child in suite.children:
-            if isinstance(child, Node) and child.type == syms.decorated:
-                for sub in child.children:
-                    if sub.type in (syms.funcdef, syms.async_funcdef):
-                        return EmptyLineTracker._get_funcdef_name(sub)
-                return None
-            if child.type in (token.NEWLINE, token.NL, token.INDENT, token.DEDENT):
-                continue
-            return None
-        return None
-
-    @staticmethod
-    def _block_is_part_of_overload_group(line: Line) -> bool:
-        """Check if a block-opening line contains a decorated function that is
-        part of a larger overload group — either because the ``if_stmt``'s next
-        sibling is a same-name decorated function, or because another branch of
-        the same ``if_stmt`` has one.
-        """
-        func_name = EmptyLineTracker._get_block_first_decorated_funcname(line)
-        if func_name is None:
-            return False
-
-        suite = line.leaves[-1].next_sibling
-        if suite is None or not isinstance(suite, Node):
-            return False
-        if_stmt = suite.parent
-        if if_stmt is None:
-            return False
-
-        # Check if the if_stmt's next sibling is a same-name decorated function.
-        sibling = if_stmt.next_sibling
-        while sibling is not None:
-            if sibling.type == syms.decorated:
-                for sub in sibling.children:
-                    if (
-                        sub.type in (syms.funcdef, syms.async_funcdef)
-                        and EmptyLineTracker._get_funcdef_name(sub) == func_name
-                    ):
-                        return True
-                break
-            elif sibling.type in (
-                token.NEWLINE,
-                token.NL,
-                token.INDENT,
-                token.DEDENT,
-                token.COMMENT,
-                token.ENDMARKER,
-            ):
-                sibling = sibling.next_sibling
-            else:
-                break
-
-        # Check other branches (elif/else) of the same if_stmt.
-        for child in if_stmt.children:
-            if (
-                isinstance(child, Node)
-                and child.type == syms.suite
-                and child is not suite
-            ):
-                for stmt in child.children:
-                    if not isinstance(stmt, Node):
-                        continue
-                    if stmt.type != syms.decorated:
-                        continue
-                    for sub in stmt.children:
-                        if (
-                            sub.type in (syms.funcdef, syms.async_funcdef)
-                            and EmptyLineTracker._get_funcdef_name(sub) == func_name
-                        ):
-                            return True
-                    break
-
-        return False
-
     def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
         """Return the number of extra empty lines before and after the `current_line`.
 
@@ -835,15 +748,7 @@ def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
                 and not current_line.is_comment
                 and (
                     self._pyi_previous_decorated_func is None
-                    or (
-                        current_line.depth <= self._pyi_previous_decorated_func.depth
-                        # Don't reset on else/elif — they continue an if/else
-                        # chain that may contain overloads at a deeper depth.
-                        and not (
-                            current_line.leaves
-                            and current_line.leaves[0].value in ("else", "elif")
-                        )
-                    )
+                    or current_line.depth <= self._pyi_previous_decorated_func.depth
                 )
             ):
                 # Only reset when we see a non-decorator line at the same or
@@ -920,32 +825,15 @@ def _maybe_empty_lines(self, current_line: Line) -> tuple[int, int]:
                     and self._pyi_previous_decorated_func.is_multi
                     and not current_line.is_comment
                     and self.previous_line.depth >= current_line.depth
-                    and not (
-                        current_line.leaves
-                        and current_line.leaves[0].value in ("else", "elif")
-                    )
                 ):
                     if self._is_in_current_group(current_line):
                         before = 0
-                    elif current_line.opens_block and (
-                        self._get_block_first_decorated_funcname(current_line)
-                        == self._pyi_previous_decorated_func.name
-                    ):
-                        before = 0
                     else:
                         before = 1
                 elif depth and not current_line.is_def and self.previous_line.is_def:
-                    if (
-                        Preview.pyi_overload_group_blank_lines in self.mode
-                        and current_line.opens_block
-                        and self.previous_line.depth <= current_line.depth
-                        and self._block_is_part_of_overload_group(current_line)
-                    ):
-                        before = 1
-                    else:
-                        # Empty lines between attributes and methods should
-                        # be preserved.
-                        before = 1 if user_had_newline else 0
+                    # Empty lines between attributes and methods should
+                    # be preserved.
+                    before = 1 if user_had_newline else 0
                 elif depth:
                     before = 0
                 else:
diff --git a/tests/data/cases/pyi_overload_groups.py b/tests/data/cases/pyi_overload_groups.py
index 2466852..1a461da 100644
--- a/tests/data/cases/pyi_overload_groups.py
+++ b/tests/data/cases/pyi_overload_groups.py
@@ -578,7 +578,6 @@ def conditional_overload(x: int) -> int:
 
 @overload
 def conditional_overload(x: str) -> str: ...
-
 def after_conditional_overload() -> None: ...
 def before_conditional_overload_with_gaps() -> None: ...
 
@@ -594,7 +593,6 @@ def conditional_overload_with_gaps(x: int) -> int:
 
 @overload
 def conditional_overload_with_gaps(x: str) -> str: ...
-
 def after_conditional_overload_with_gaps() -> None: ...
 
 if sys.version_info >= (3, 10):
@@ -663,7 +661,6 @@ def after_comment_group() -> None: ...
 
 class ClassWithConditionalOverloads:
     def method(self) -> None: ...
-
     if sys.version_info >= (3, 10):
         @overload
         def conditional(self, x: int) -> int: ...
@@ -679,6 +676,7 @@ class ClassWithOverloadThenConditional:
     def method(self, x: int) -> int: ...
     @overload
     def method(self, x: str) -> str: ...
+
     if sys.version_info >= (3, 10):
         @overload
         def method(self, x: bytes) -> bytes: ...

[srittau]

This new formatting makes sense to me. Decorators are a "visual break", separating two overloaded functions from each other in a way. Especially the last overload looks more tightly coupled to the succeeding function than the previous overload. This will hopefully help to group overloads more closely.

[Akuli]

I like the new formatting too.

[MeGaGiGaGon]

I'm surprised at how much code this takes, but the results seem worth it. The behavior in the new test looks good to me, as well as the new logic, though I didn't trace it fully. I assume we've just been waiting on feedback in python/typeshed#15504 before merging this.