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

CHANGES.md

@@ -20,6 +20,9 @@
 
 <!-- Changes that affect Black's preview style -->
 
+- Improve heuristics around whether blank lines should appear before, within and after
+  groups of same-name decorated functions (such as `@overload` groups) in `.pyi` stub
+  files (#5021)
 - Fix bug where `if` guards in `case` blocks were incorrectly split when the pattern had
   a trailing comma (#4884)
 - Fix `string_processing` crashing on unassigned long string literals with trailing

docs/the_black_code_style/future_style.md

@@ -25,6 +25,9 @@ Currently, the following features are included in the preview style:
 - `fix_if_guard_explosion_in_case_statement`: fixed exploding of the if guard in case
   patterns which have trailing commas in them, even if the guard expression fits in one
   line
+- `pyi_overload_group_blank_lines`: In `.pyi` stub files, improve heuristics around when
+  blank lines should appear before, after and within decorated function groups.
+  ([see below](labels/pyi-overload-group))
 
 (labels/wrap-comprehension-in)=
 
@@ -137,6 +140,55 @@ my_dict = {
 }
 ```
 
+(labels/pyi-overload-group)=
+
+### Improved overload groups in stub files
+
+In `.pyi` stub files, Black now has improved heuristics regarding when blank lines
+should appear before, after or within groups of decorated functions that share the same
+name (such as `@overload` groups). Two rules are applied when a decorated function is
+determined to be part of a series of >=2 decorated functions with the same name:
+
+1. **Before the decorated function**: a blank line is always inserted, unless the
+   preceding statement is a same-name decorated function (i.e. part of an `@overload`
+   group) or the function is the first statement in its block.
+2. **After the decorated function**: a blank line is always inserted, unless the
+   following statement is a same-name decorated function.
+
+These rules apply regardless of what the adjacent statement is — whether it's another
+function definition, a variable annotation, or any other statement.
+
+Previously, Black could insert unwanted blank lines _within_ an overload group when one
+of the overloads had a docstring, and did not consistently enforce blank lines at the
+boundaries of overload groups:
+
+```python
+# Before
+
+@overload
+def foo(x: int) -> int:
+    """Docs."""
+
+@overload                    # unwanted blank line within group
+def foo(x: str) -> str: ...
+def bar(x): ...              # no blank line after group
+```
+
+With this feature enabled, the group is kept together and clearly separated from
+surrounding code:
+
+```python
+# After (with --preview)
+
+@overload
+def foo(x: int) -> int:
+    """Docs."""
+@overload
+def foo(x: str) -> str: ...
+
+def bar(x): ...
+```
+
 ## Unstable style
 
 (labels/unstable-style)=

src/black/lines.py

@@ -5,7 +5,7 @@
 from typing import Optional, TypeVar, Union, cast
 
 from black.brackets import COMMA_PRIORITY, DOT_PRIORITY, BracketTracker
-from black.mode import Mode
+from black.mode import Mode, Preview
 from black.nodes import (
     BRACKETS,
     CLOSING_BRACKETS,
@@ -544,6 +544,132 @@ class EmptyLineTracker:
     previous_block: LinesBlock | None = None
     previous_defs: list[Line] = field(default_factory=list)
     semantic_leading_comment: LinesBlock | None = None
+    _pyi_previous_decorated_func: tuple[str, int, bool] | None = None
+
+    @staticmethod
+    def _get_funcdef_name(node: Node | Leaf) -> str | None:
+        """Extract the function name from a funcdef or async_funcdef node."""
+        funcdef: Node | Leaf = node
+        if isinstance(node, Node) and node.type == syms.async_funcdef:
+            for sub in node.children:
+                if sub.type == syms.funcdef:
+                    funcdef = sub
+                    break
+        if not isinstance(funcdef, Node) or funcdef.type != syms.funcdef:
+            return None
+        for child in funcdef.children:
+            if (
+                isinstance(child, Leaf)
+                and child.type == token.NAME
+                and child.value != "def"
+            ):
+                return child.value
+        return None
+
+    @staticmethod
+    def _get_def_name(line: Line) -> str | None:
+        """Extract the function name from a line that is a function definition."""
+        if not line.is_def or not line.leaves:
+            return None
+        if line.leaves[0].value == "def":
+            return line.leaves[1].value
+        # async def
+        return line.leaves[2].value
+
+    @staticmethod
+    def _is_line_decorated(line: Line) -> bool:
+        """Check if a def line is part of a decorated statement."""
+        if not line.is_def or not line.leaves:
+            return False
+        node = line.leaves[0].parent
+        while node is not None:
+            if node.type == syms.decorated:
+                return True
+            # Stop at class/function boundaries to avoid matching a
+            # decorated *class* that contains this function.
+            if node.type in (syms.classdef, syms.funcdef, syms.async_funcdef):
+                node = node.parent
+                continue
+            if node.type == syms.suite:
+                # A suite means we've entered a class/function body;
+                # any decorated node above here belongs to the enclosing
+                # scope, not to this function.
+                return False
+            node = node.parent
+        return False
+
+    @staticmethod
+    def _find_decorated_node(line: Line) -> Node | None:
+        """Walk up from a decorator line's '@' leaf to its `decorated` node."""
+        if not line.leaves:
+            return None
+        node = line.leaves[0].parent
+        while node is not None:
+            if node.type == syms.decorated:
+                assert isinstance(node, Node)
+                return node
+            if node.type in (syms.suite, syms.file_input):
+                return None
+            node = node.parent
+        return None
+
+    @staticmethod
+    def _get_decorator_target_name(line: Line) -> str | None:
+        """For a decorator line, extract the name of the function it decorates.
+
+        Only handles decorated functions, not decorated classes.
+        """
+        if not line.is_decorator:
+            return None
+        decorated = EmptyLineTracker._find_decorated_node(line)
+        if decorated is None:
+            return None
+        for child in decorated.children:
+            if child.type in (syms.funcdef, syms.async_funcdef):
+                return EmptyLineTracker._get_funcdef_name(child)
+        return None
+
+    @staticmethod
+    def _is_start_of_decorated_group(line: Line) -> bool:
+        """Check if a decorator line starts a multi-function group.
+
+        A multi-function group is 2+ consecutive decorated functions sharing the
+        same name (e.g. @overload groups, @property + setter pairs).
+        Returns True when the very next statement-level sibling of the current
+        decorated node is also a decorated function with the same name.
+        """
+        if not line.is_decorator:
+            return False
+
+        cur_name = EmptyLineTracker._get_decorator_target_name(line)
+        if cur_name is None:
+            return False
+
+        decorated_node = EmptyLineTracker._find_decorated_node(line)
+        if decorated_node is None or decorated_node.parent is None:
+            return False
+
+        # Scan forward through siblings to find the next statement node.
+        sibling = decorated_node.next_sibling
+        while sibling is not None:
+            if sibling.type == syms.decorated:
+                for subchild in sibling.children:
+                    if subchild.type in (syms.funcdef, syms.async_funcdef):
+                        return EmptyLineTracker._get_funcdef_name(subchild) == cur_name
+                return False
+            elif sibling.type in (
+                token.NEWLINE,
+                token.NL,
+                token.INDENT,
+                token.DEDENT,
+                token.COMMENT,
+                token.ENDMARKER,
+            ):
+                sibling = sibling.next_sibling
+            else:
+                return False
+
+        return False
 
     def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
         """Return the number of extra empty lines before and after the `current_line`.
@@ -586,6 +712,45 @@ def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
         elif not current_line.is_decorator or before:
             self.semantic_leading_comment = None
 
+        # Maintain _pyi_previous_decorated_func state for overload groups.
+        # The tuple is (name, depth, is_multi) where is_multi indicates
+        # the group has 2+ same-name decorated functions.
+        if self.mode.is_pyi and Preview.pyi_overload_group_blank_lines in self.mode:
+            if current_line.is_def and self._is_line_decorated(current_line):
+                name = self._get_def_name(current_line)
+                if name is not None:
+                    prev = self._pyi_previous_decorated_func
+                    is_multi = (
+                        prev is not None
+                        and prev[0] == name
+                        and prev[1] == current_line.depth
+                    )
+                    self._pyi_previous_decorated_func = (
+                        name,
+                        current_line.depth,
+                        is_multi or (prev is not None and prev[0] == name and prev[2]),
+                    )
+            elif (
+                not current_line.is_decorator
+                and not current_line.is_comment
+                and (
+                    self._pyi_previous_decorated_func is None
+                    or (
+                        current_line.depth <= self._pyi_previous_decorated_func[1]
+                        # 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")
+                        )
+                    )
+                )
+            ):
+                # Only reset when we see a non-decorator line at the same or
+                # lower depth.  Body lines (docstrings, ...) at deeper depth
+                # should not clear the state.
+                self._pyi_previous_decorated_func = None
+
         self.previous_line = current_line
         self.previous_block = block
         return block
@@ -642,9 +807,36 @@ def _maybe_empty_lines(self, current_line: Line) -> tuple[int, int]:
 
         if previous_def is not None:
             assert self.previous_line is not None
+            # Note: for decorator/def/class lines, `before` computed here is
+            # passed to _maybe_empty_lines_for_class_or_def which may override
+            # it. This block still matters for non-decorator/def/class lines
+            # (e.g. a `var: int` statement following an overload group).
             if self.mode.is_pyi:
                 if previous_def.is_class and not previous_def.is_stub_class:
                     before = 1
+                elif (
+                    Preview.pyi_overload_group_blank_lines in self.mode
+                    and self._pyi_previous_decorated_func is not None
+                    and self._pyi_previous_decorated_func[2]
+                ):
+                    prev_name, prev_depth, _ = self._pyi_previous_decorated_func
+                    cur_name = (
+                        self._get_decorator_target_name(current_line)
+                        if current_line.is_decorator
+                        else self._get_def_name(current_line)
+                    )
+                    if (
+                        cur_name is not None
+                        and cur_name == prev_name
+                        and prev_depth == depth
+                        and (
+                            current_line.is_decorator
+                            or self._is_line_decorated(current_line)
+                        )
+                    ):
+                        before = 0
+                    else:
+                        before = 1
                 elif depth and not current_line.is_def and self.previous_line.is_def:
                     # Empty lines between attributes and methods should be preserved.
                     before = 1 if user_had_newline else 0
@@ -750,6 +942,46 @@ def _maybe_empty_lines_for_class_or_def(
             # statement in the same level, we always want a blank line if there's
             # something with a body preceding.
             elif self.previous_line.depth > current_line.depth:
+                if (
+                    Preview.pyi_overload_group_blank_lines in self.mode
+                    and self._pyi_previous_decorated_func is not None
+                    and current_line.is_decorator
+                ):
+                    prev_name, prev_depth, _ = self._pyi_previous_decorated_func
+                    cur_name = self._get_decorator_target_name(current_line)
+                    if cur_name == prev_name and prev_depth == current_line.depth:
+                        newlines = 0
+                    else:
+                        newlines = 1
+                else:
+                    newlines = 1
+            elif (
+                Preview.pyi_overload_group_blank_lines in self.mode
+                and self._pyi_previous_decorated_func is not None
+                and self._pyi_previous_decorated_func[2]
+            ):
+                prev_name, prev_depth, _ = self._pyi_previous_decorated_func
+                cur_name = (
+                    self._get_decorator_target_name(current_line)
+                    if current_line.is_decorator
+                    else self._get_def_name(current_line)
+                )
+                is_same_group = (
+                    cur_name is not None
+                    and cur_name == prev_name
+                    and prev_depth == current_line.depth
+                    and (
+                        current_line.is_decorator
+                        or self._is_line_decorated(current_line)
+                    )
+                )
+                newlines = 0 if is_same_group else 1
+            elif (
+                Preview.pyi_overload_group_blank_lines in self.mode
+                and current_line.is_decorator
+                and self.previous_line.depth >= current_line.depth
+                and self._is_start_of_decorated_group(current_line)
+            ):
                 newlines = 1
             elif (
                 current_line.is_def or current_line.is_decorator

src/black/mode.py

@@ -230,6 +230,7 @@ class Preview(Enum):
     simplify_power_operator_hugging = auto()
     wrap_long_dict_values_in_parens = auto()
     fix_if_guard_explosion_in_case_statement = auto()
+    pyi_overload_group_blank_lines = auto()
 
 
 UNSTABLE_FEATURES: set[Preview] = {

src/black/resources/black.schema.json

@@ -85,7 +85,8 @@
           "wrap_comprehension_in",
           "simplify_power_operator_hugging",
           "wrap_long_dict_values_in_parens",
-          "fix_if_guard_explosion_in_case_statement"
+          "fix_if_guard_explosion_in_case_statement",
+          "pyi_overload_group_blank_lines"
         ]
       },
       "description": "Enable specific features included in the `--unstable` style. Requires `--preview`. No compatibility guarantees are provided on the behavior or existence of any unstable features."