From 8c5d96ffd3da6d621e67639dadd26a1d7d0227cd Mon Sep 17 00:00:00 2001
From: John Litborn <11260241+jakkdl@users.noreply.github.com>
Date: Fri, 22 Sep 2023 17:38:51 +0200
Subject: [PATCH] fix indentation of line breaks in long type hints by adding
 parens (#3899)

* fix indentation of line breaks in long type hints by adding parentheses, and remove unnecessary parentheses

* add entry in CHANGES.md, make the style change only in preview mode
---
 CHANGES.md                                    |   3 +
 src/black/linegen.py                          |  30 ++-
 src/black/mode.py                             |   1 +
 .../preview/long_strings__type_annotations.py |   2 +-
 .../pep604_union_types_line_breaks.py         | 187 ++++++++++++++++++
 5 files changed, 220 insertions(+), 3 deletions(-)
 create mode 100644 tests/data/preview_py_310/pep604_union_types_line_breaks.py

diff --git a/CHANGES.md b/CHANGES.md
index a68106ad23f..a879ab3e8da 100644
--- a/CHANGES.md
+++ b/CHANGES.md
@@ -12,6 +12,9 @@
 
 ### Preview style
 
+- Long type hints are now wrapped in parentheses and properly indented when split across
+  multiple lines (#3899)
+
 <!-- Changes that affect Black's preview style -->
 
 ### Configuration
diff --git a/src/black/linegen.py b/src/black/linegen.py
index 507e860190f..9ddd4619f69 100644
--- a/src/black/linegen.py
+++ b/src/black/linegen.py
@@ -397,6 +397,24 @@ def visit_factor(self, node: Node) -> Iterator[Line]:
             node.insert_child(index, Node(syms.atom, [lpar, operand, rpar]))
         yield from self.visit_default(node)
 
+    def visit_tname(self, node: Node) -> Iterator[Line]:
+        """
+        Add potential parentheses around types in function parameter lists to be made
+        into real parentheses in case the type hint is too long to fit on a line
+        Examples:
+        def foo(a: int, b: float = 7): ...
+
+        ->
+
+        def foo(a: (int), b: (float) = 7): ...
+        """
+        if Preview.parenthesize_long_type_hints in self.mode:
+            assert len(node.children) == 3
+            if maybe_make_parens_invisible_in_atom(node.children[2], parent=node):
+                wrap_in_parentheses(node, node.children[2], visible=False)
+
+        yield from self.visit_default(node)
+
     def visit_STRING(self, leaf: Leaf) -> Iterator[Line]:
         if Preview.hex_codes_in_unicode_sequences in self.mode:
             normalize_unicode_escape_sequences(leaf)
@@ -498,7 +516,14 @@ def __post_init__(self) -> None:
         self.visit_except_clause = partial(v, keywords={"except"}, parens={"except"})
         self.visit_with_stmt = partial(v, keywords={"with"}, parens={"with"})
         self.visit_classdef = partial(v, keywords={"class"}, parens=Ø)
-        self.visit_expr_stmt = partial(v, keywords=Ø, parens=ASSIGNMENTS)
+
+        # When this is moved out of preview, add ":" directly to ASSIGNMENTS in nodes.py
+        if Preview.parenthesize_long_type_hints in self.mode:
+            assignments = ASSIGNMENTS | {":"}
+        else:
+            assignments = ASSIGNMENTS
+        self.visit_expr_stmt = partial(v, keywords=Ø, parens=assignments)
+
         self.visit_return_stmt = partial(v, keywords={"return"}, parens={"return"})
         self.visit_import_from = partial(v, keywords=Ø, parens={"import"})
         self.visit_del_stmt = partial(v, keywords=Ø, parens={"del"})
@@ -1368,7 +1393,7 @@ def maybe_make_parens_invisible_in_atom(
     Returns whether the node should itself be wrapped in invisible parentheses.
     """
     if (
-        node.type != syms.atom
+        node.type not in (syms.atom, syms.expr)
         or is_empty_tuple(node)
         or is_one_tuple(node)
         or (is_yield(node) and parent.type != syms.expr_stmt)
@@ -1392,6 +1417,7 @@ def maybe_make_parens_invisible_in_atom(
             syms.except_clause,
             syms.funcdef,
             syms.with_stmt,
+            syms.tname,
             # these ones aren't useful to end users, but they do please fuzzers
             syms.for_stmt,
             syms.del_stmt,
diff --git a/src/black/mode.py b/src/black/mode.py
index 8a855ac495a..f44a821bcd0 100644
--- a/src/black/mode.py
+++ b/src/black/mode.py
@@ -180,6 +180,7 @@ class Preview(Enum):
     # for https://github.com/psf/black/issues/3117 to be fixed.
     string_processing = auto()
     parenthesize_conditional_expressions = auto()
+    parenthesize_long_type_hints = auto()
     skip_magic_trailing_comma_in_subscript = auto()
     wrap_long_dict_values_in_parens = auto()
     wrap_multiple_context_managers_in_parens = auto()
diff --git a/tests/data/preview/long_strings__type_annotations.py b/tests/data/preview/long_strings__type_annotations.py
index 41d7ee2b67b..45de882d02c 100644
--- a/tests/data/preview/long_strings__type_annotations.py
+++ b/tests/data/preview/long_strings__type_annotations.py
@@ -54,6 +54,6 @@ def func(
 
 
 def func(
-    argument: ("int |" "str"),
+    argument: "int |" "str",
 ) -> Set["int |" " str"]:
     pass
diff --git a/tests/data/preview_py_310/pep604_union_types_line_breaks.py b/tests/data/preview_py_310/pep604_union_types_line_breaks.py
new file mode 100644
index 00000000000..9c4ab870766
--- /dev/null
+++ b/tests/data/preview_py_310/pep604_union_types_line_breaks.py
@@ -0,0 +1,187 @@
+# This has always worked
+z= Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong
+
+# "AnnAssign"s now also work
+z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong
+z: (Short
+    | Short2
+    | Short3
+    | Short4)
+z: (int)
+z: ((int))
+
+
+z: Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong | Loooooooooooooooooooooooong = 7
+z: (Short
+    | Short2
+    | Short3
+    | Short4) = 8
+z: (int) = 2.3
+z: ((int)) = foo()
+
+# In case I go for not enforcing parantheses, this might get improved at the same time
+x = (
+    z
+    == 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999,
+    y
+    == 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999,
+)
+
+x = (
+    z == (9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999),
+    y == (9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999),
+)
+
+# handle formatting of "tname"s in parameter list
+
+# remove unnecessary paren
+def foo(i: (int)) -> None: ...
+
+
+# this is a syntax error in the type annotation according to mypy, but it's not invalid *python* code, so make sure we don't mess with it and make it so.
+def foo(i: (int,)) -> None: ...
+
+def foo(
+    i: int,
+    x: Loooooooooooooooooooooooong
+    | Looooooooooooooooong
+    | Looooooooooooooooooooong
+    | Looooooong,
+    *,
+    s: str,
+) -> None:
+    pass
+
+
+@app.get("/path/")
+async def foo(
+    q: str
+    | None = Query(None, title="Some long title", description="Some long description")
+):
+    pass
+
+
+def f(
+    max_jobs: int
+    | None = Option(
+        None, help="Maximum number of jobs to launch. And some additional text."
+        ),
+    another_option: bool = False
+    ):
+    ...
+
+
+# output
+# This has always worked
+z = (
+    Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+)
+
+# "AnnAssign"s now also work
+z: (
+    Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+)
+z: Short | Short2 | Short3 | Short4
+z: int
+z: int
+
+
+z: (
+    Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+    | Loooooooooooooooooooooooong
+) = 7
+z: Short | Short2 | Short3 | Short4 = 8
+z: int = 2.3
+z: int = foo()
+
+# In case I go for not enforcing parantheses, this might get improved at the same time
+x = (
+    z
+    == 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999
+    | 9999999999999999999999999999999999999999,
+    y
+    == 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999
+    + 9999999999999999999999999999999999999999,
+)
+
+x = (
+    z
+    == (
+        9999999999999999999999999999999999999999
+        | 9999999999999999999999999999999999999999
+        | 9999999999999999999999999999999999999999
+        | 9999999999999999999999999999999999999999
+    ),
+    y
+    == (
+        9999999999999999999999999999999999999999
+        + 9999999999999999999999999999999999999999
+        + 9999999999999999999999999999999999999999
+        + 9999999999999999999999999999999999999999
+    ),
+)
+
+# handle formatting of "tname"s in parameter list
+
+
+# remove unnecessary paren
+def foo(i: int) -> None: ...
+
+
+# this is a syntax error in the type annotation according to mypy, but it's not invalid *python* code, so make sure we don't mess with it and make it so.
+def foo(i: (int,)) -> None: ...
+
+
+def foo(
+    i: int,
+    x: (
+        Loooooooooooooooooooooooong
+        | Looooooooooooooooong
+        | Looooooooooooooooooooong
+        | Looooooong
+    ),
+    *,
+    s: str,
+) -> None:
+    pass
+
+
+@app.get("/path/")
+async def foo(
+    q: str | None = Query(
+        None, title="Some long title", description="Some long description"
+    )
+):
+    pass
+
+
+def f(
+    max_jobs: int | None = Option(
+        None, help="Maximum number of jobs to launch. And some additional text."
+    ),
+    another_option: bool = False,
+): ...