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,57 @@ 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/pyi-overload-group)=
+
+### Improved heuristics for blank lines before, after and within decorated function 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:
+
+1. **Before a 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 a 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): ...
+```
 
 (labels/wrap-comprehension-in)=
 

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,69 @@ 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] | None = 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:
+            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 _get_decorator_target_name(line: Line) -> str | None:
+        """For a decorator line, extract the name of the function it decorates."""
+        if not line.is_decorator or not line.leaves:
+            return None
+        node = line.leaves[0].parent
+        while node is not None:
+            if node.type != syms.decorated:
+                node = node.parent
+                continue
+            for child in node.children:
+                funcdef = child
+                if child.type == syms.async_funcdef:
+                    # async_funcdef has children: ASYNC, funcdef
+                    for sub in child.children:
+                        if sub.type == syms.funcdef:
+                            funcdef = sub
+                            break
+                if funcdef.type != syms.funcdef:
+                    continue
+                for leaf in funcdef.children:
+                    if (
+                        isinstance(leaf, Leaf)
+                        and leaf.type == token.NAME
+                        and leaf.value != "def"
+                    ):
+                        return leaf.value
+            break
+        return None
 
     def maybe_empty_lines(self, current_line: Line) -> LinesBlock:
         """Return the number of extra empty lines before and after the `current_line`.
@@ -586,6 +649,21 @@ 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.
+        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:
+                    self._pyi_previous_decorated_func = (name, current_line.depth)
+            elif not current_line.is_decorator and (
+                self._pyi_previous_decorated_func is None
+                or current_line.depth <= self._pyi_previous_decorated_func[1]
+            ):
+                # 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
@@ -645,6 +723,28 @@ def _maybe_empty_lines(self, current_line: Line) -> tuple[int, int]:
             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
+                ):
+                    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 +850,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
+            ):
+                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
+            ):
+                # Before a decorated function, enforce blank line unless
+                # first statement in the block.
                 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."