Improve function declaration wrapping when it contains generic type definitions (#4553)

---------

Co-authored-by: Jelle Zijlstra <jelle.zijlstra@gmail.com>
Co-authored-by: hauntsaninja <hauntsaninja@gmail.com>
Co-authored-by: Shantanu <12621235+hauntsaninja@users.noreply.github.com>

CHANGES.md

@@ -22,6 +22,8 @@ the following changes:
 The following changes were not in any previous release:
 
 - Remove parentheses around sole list items (#4312)
+- Generic function definitions are now formatted more elegantly: parameters are
+  split over multiple lines first instead of type parameter definitions (#4553)
 
 ### Stable style
 

src/black/linegen.py

@@ -779,26 +779,29 @@ def left_hand_split(
     Prefer RHS otherwise.  This is why this function is not symmetrical with
     :func:`right_hand_split` which also handles optional parentheses.
     """
-    tail_leaves: list[Leaf] = []
-    body_leaves: list[Leaf] = []
-    head_leaves: list[Leaf] = []
-    current_leaves = head_leaves
-    matching_bracket: Optional[Leaf] = None
-    for leaf in line.leaves:
-        if (
-            current_leaves is body_leaves
-            and leaf.type in CLOSING_BRACKETS
-            and leaf.opening_bracket is matching_bracket
-            and isinstance(matching_bracket, Leaf)
-        ):
-            ensure_visible(leaf)
-            ensure_visible(matching_bracket)
-            current_leaves = tail_leaves if body_leaves else head_leaves
-        current_leaves.append(leaf)
-        if current_leaves is head_leaves:
-            if leaf.type in OPENING_BRACKETS:
-                matching_bracket = leaf
-                current_leaves = body_leaves
+    for leaf_type in [token.LPAR, token.LSQB]:
+        tail_leaves: list[Leaf] = []
+        body_leaves: list[Leaf] = []
+        head_leaves: list[Leaf] = []
+        current_leaves = head_leaves
+        matching_bracket: Optional[Leaf] = None
+        for leaf in line.leaves:
+            if (
+                current_leaves is body_leaves
+                and leaf.type in CLOSING_BRACKETS
+                and leaf.opening_bracket is matching_bracket
+                and isinstance(matching_bracket, Leaf)
+            ):
+                ensure_visible(leaf)
+                ensure_visible(matching_bracket)
+                current_leaves = tail_leaves if body_leaves else head_leaves
+            current_leaves.append(leaf)
+            if current_leaves is head_leaves:
+                if leaf.type == leaf_type:
+                    matching_bracket = leaf
+                    current_leaves = body_leaves
+        if matching_bracket and tail_leaves:
+            break
     if not matching_bracket or not tail_leaves:
         raise CannotSplit("No brackets found")
 

tests/data/cases/generics_wrapping.py

@@ -0,0 +1,307 @@
+# flags: --minimum-version=3.12
+def plain[T, B](a: T, b: T) -> T:
+    return a
+
+def arg_magic[T, B](a: T, b: T,) -> T:
+    return a
+
+def type_param_magic[T, B,](a: T, b: T) -> T:
+    return a
+
+def both_magic[T, B,](a: T, b: T,) -> T:
+    return a
+
+
+def plain_multiline[
+    T,
+    B
+](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def arg_magic_multiline[
+    T,
+    B
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+def type_param_magic_multiline[
+    T,
+    B,
+](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def both_magic_multiline[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def plain_mixed1[
+    T,
+    B
+](a: T, b: T) -> T:
+    return a
+
+def plain_mixed2[T, B](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def arg_magic_mixed1[
+    T,
+    B
+](a: T, b: T,) -> T:
+    return a
+
+def arg_magic_mixed2[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+def type_param_magic_mixed1[
+    T,
+    B,
+](a: T, b: T) -> T:
+    return a
+
+def type_param_magic_mixed2[T, B,](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def both_magic_mixed1[
+    T,
+    B,
+](a: T, b: T,) -> T:
+    return a
+
+def both_magic_mixed2[T, B,](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+def something_something_function[
+    T: Model
+](param: list[int], other_param: type[T], *, some_other_param: bool = True) -> QuerySet[
+    T
+]:
+    pass
+
+
+def func[A_LOT_OF_GENERIC_TYPES: AreBeingDefinedHere, LIKE_THIS, AND_THIS, ANOTHER_ONE, AND_YET_ANOTHER_ONE: ThisOneHasTyping](a: T, b: T, c: T, d: T, e: T, f: T, g: T, h: T, i: T, j: T, k: T, l: T, m: T, n: T, o: T, p: T) -> T:
+    return a
+
+
+def with_random_comments[
+    Z
+    # bye
+]():
+    return a
+
+
+def func[
+    T,  # comment
+    U  # comment
+    ,
+    Z:  # comment
+    int
+](): pass
+
+
+def func[
+    T,  # comment but it's long so it doesn't just move to the end of the line
+    U  # comment comment comm comm ent ent
+    ,
+    Z:  # comment ent ent comm comm comment
+    int
+](): pass
+
+
+# output
+def plain[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def type_param_magic[
+    T,
+    B,
+](
+    a: T, b: T
+) -> T:
+    return a
+
+
+def both_magic[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def plain_multiline[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic_multiline[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def type_param_magic_multiline[
+    T,
+    B,
+](
+    a: T, b: T
+) -> T:
+    return a
+
+
+def both_magic_multiline[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def plain_mixed1[T, B](a: T, b: T) -> T:
+    return a
+
+
+def plain_mixed2[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic_mixed1[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def arg_magic_mixed2[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def type_param_magic_mixed1[
+    T,
+    B,
+](
+    a: T, b: T
+) -> T:
+    return a
+
+
+def type_param_magic_mixed2[
+    T,
+    B,
+](
+    a: T, b: T
+) -> T:
+    return a
+
+
+def both_magic_mixed1[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def both_magic_mixed2[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def something_something_function[T: Model](
+    param: list[int], other_param: type[T], *, some_other_param: bool = True
+) -> QuerySet[T]:
+    pass
+
+
+def func[
+    A_LOT_OF_GENERIC_TYPES: AreBeingDefinedHere,
+    LIKE_THIS,
+    AND_THIS,
+    ANOTHER_ONE,
+    AND_YET_ANOTHER_ONE: ThisOneHasTyping,
+](
+    a: T,
+    b: T,
+    c: T,
+    d: T,
+    e: T,
+    f: T,
+    g: T,
+    h: T,
+    i: T,
+    j: T,
+    k: T,
+    l: T,
+    m: T,
+    n: T,
+    o: T,
+    p: T,
+) -> T:
+    return a
+
+
+def with_random_comments[
+    Z
+    # bye
+]():
+    return a
+
+
+def func[T, U, Z: int]():  # comment  # comment  # comment
+    pass
+
+
+def func[
+    T,  # comment but it's long so it doesn't just move to the end of the line
+    U,  # comment comment comm comm ent ent
+    Z: int,  # comment ent ent comm comm comment
+]():
+    pass

tests/data/cases/skip_magic_trailing_comma_generic_wrap.py

@@ -0,0 +1,163 @@
+# flags: --minimum-version=3.12 --skip-magic-trailing-comma
+def plain[T, B](a: T, b: T) -> T:
+    return a
+
+def arg_magic[T, B](a: T, b: T,) -> T:
+    return a
+
+def type_param_magic[T, B,](a: T, b: T) -> T:
+    return a
+
+def both_magic[T, B,](a: T, b: T,) -> T:
+    return a
+
+
+def plain_multiline[
+    T,
+    B
+](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def arg_magic_multiline[
+    T,
+    B
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+def type_param_magic_multiline[
+    T,
+    B,
+](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def both_magic_multiline[
+    T,
+    B,
+](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+def plain_mixed1[
+    T,
+    B
+](a: T, b: T) -> T:
+    return a
+
+def plain_mixed2[T, B](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def arg_magic_mixed1[
+    T,
+    B
+](a: T, b: T,) -> T:
+    return a
+
+def arg_magic_mixed2[T, B](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+def type_param_magic_mixed1[
+    T,
+    B,
+](a: T, b: T) -> T:
+    return a
+
+def type_param_magic_mixed2[T, B,](
+    a: T,
+    b: T
+) -> T:
+    return a
+
+def both_magic_mixed1[
+    T,
+    B,
+](a: T, b: T,) -> T:
+    return a
+
+def both_magic_mixed2[T, B,](
+    a: T,
+    b: T,
+) -> T:
+    return a
+
+
+# output
+def plain[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic[T, B](a: T, b: T) -> T:
+    return a
+
+
+def type_param_magic[T, B](a: T, b: T) -> T:
+    return a
+
+
+def both_magic[T, B](a: T, b: T) -> T:
+    return a
+
+
+def plain_multiline[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic_multiline[T, B](a: T, b: T) -> T:
+    return a
+
+
+def type_param_magic_multiline[T, B](a: T, b: T) -> T:
+    return a
+
+
+def both_magic_multiline[T, B](a: T, b: T) -> T:
+    return a
+
+
+def plain_mixed1[T, B](a: T, b: T) -> T:
+    return a
+
+
+def plain_mixed2[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic_mixed1[T, B](a: T, b: T) -> T:
+    return a
+
+
+def arg_magic_mixed2[T, B](a: T, b: T) -> T:
+    return a
+
+
+def type_param_magic_mixed1[T, B](a: T, b: T) -> T:
+    return a
+
+
+def type_param_magic_mixed2[T, B](a: T, b: T) -> T:
+    return a
+
+
+def both_magic_mixed1[T, B](a: T, b: T) -> T:
+    return a
+
+
+def both_magic_mixed2[T, B](a: T, b: T) -> T:
+    return a

tests/data/cases/type_param_defaults.py

@@ -37,25 +37,27 @@ def trailing_comma2[T=int](a: str,):
 ] = something_that_is_long
 
 
-def simple[
-    T = something_that_is_long
-](short1: int, short2: str, short3: bytes) -> float:
+def simple[T = something_that_is_long](
+    short1: int, short2: str, short3: bytes
+) -> float:
     pass
 
 
-def longer[
-    something_that_is_long = something_that_is_long
-](something_that_is_long: something_that_is_long) -> something_that_is_long:
+def longer[something_that_is_long = something_that_is_long](
+    something_that_is_long: something_that_is_long,
+) -> something_that_is_long:
     pass
 
 
 def trailing_comma1[
     T = int,
-](a: str):
+](
+    a: str,
+):
     pass
 
 
-def trailing_comma2[
-    T = int
-](a: str,):
+def trailing_comma2[T = int](
+    a: str,
+):
     pass