Handle multi-line annotations correctly (#118)

Author: vrslev

auto_typing_final/transform.py

@@ -115,9 +115,9 @@ def _match_exact_identifier(node: SgNode, imports_result: ImportsResult, identif
     return None
 
 
-def _strip_identifier_from_type_annotation(  # noqa: C901, PLR0911
+def _strip_value_from_type_annotation_that_is_indeed_inside_given_identifier(
     node: SgNode, imports_result: ImportsResult, identifier_name: str
-) -> tuple[str, str] | None:
+) -> str | None:
     type_node_children: Final = node.children()
     if len(type_node_children) != 1:
         return None
@@ -127,55 +127,54 @@ def _strip_identifier_from_type_annotation(  # noqa: C901, PLR0911
     if kind == "subscript":
         match tuple((child.kind(), child) for child in inner_type_node.children()):
             case (("attribute", attribute), ("[", _), *kinds_and_nodes, ("]", _)):
-                if existing_final := _match_exact_identifier(attribute, imports_result, identifier_name):
-                    return existing_final, "".join(node.text() for _, node in kinds_and_nodes)
+                if _match_exact_identifier(attribute, imports_result, identifier_name):
+                    return "".join(node.text() for _, node in kinds_and_nodes)
     elif kind == "generic_type" and imports_result.has_from_import:
         match tuple((child.kind(), child) for child in inner_type_node.children()):
             case (("identifier", identifier), ("type_parameter", type_parameter)):
                 if identifier.text() != identifier_name:
                     return None
                 match tuple((inner_child.kind(), inner_child) for inner_child in type_parameter.children()):
                     case (("[", _), *kinds_and_nodes, ("]", _)):
-                        return identifier_name, "".join(node.text() for _, node in kinds_and_nodes)
-    elif kind == "identifier" and inner_type_node.text() == identifier_name:
-        return identifier_name, ""
-    elif kind == "attribute" and (
-        existing_final := _match_exact_identifier(inner_type_node, imports_result, identifier_name)
+                        return "".join(node.text() for _, node in kinds_and_nodes)
+    elif (kind == "identifier" and inner_type_node.text() == identifier_name) or (
+        kind == "attribute" and _match_exact_identifier(inner_type_node, imports_result, identifier_name)
     ):
-        return existing_final, ""
+        return ""
     return None
 
 
-def _make_changed_text_from_operation(  # noqa: C901
+def _make_changed_text_from_operation(
     operation: Operation, final_value: str, imports_result: ImportsResult, identifier_name: str
 ) -> Iterable[tuple[SgNode, str]]:
     match operation:
         case AddFinal(assignment):
             match assignment:
                 case EditableAssignmentWithoutAnnotation(node, left, right):
                     yield node, f"{left}: {final_value} = {right}"
-                case EditableAssignmentWithAnnotation(node, left, annotation, right):
-                    match _strip_identifier_from_type_annotation(annotation, imports_result, identifier_name):
-                        case None:
-                            yield node, f"{left}: {final_value}[{annotation.text()}] = {right}"
-                        case existing_final, "":
-                            yield node, f"{left}: {existing_final} = {right}"
-                        case existing_final, new_annotation:
-                            yield node, f"{left}: {existing_final}[{new_annotation}] = {right}"
-
+                case EditableAssignmentWithAnnotation(node, left, annotation, right) if (
+                    _strip_value_from_type_annotation_that_is_indeed_inside_given_identifier(
+                        annotation, imports_result, identifier_name
+                    )
+                    is None
+                ):
+                    yield node, f"{left}: {final_value}[{annotation.text()}] = {right}"
         case RemoveFinal(assignments):
             for assignment in assignments:
                 match assignment:
                     case EditableAssignmentWithoutAnnotation(node, left, right):
                         yield node, node.text()
                     case EditableAssignmentWithAnnotation(node, left, annotation, right):
-                        match _strip_identifier_from_type_annotation(annotation, imports_result, identifier_name):
-                            case _, "":
+                        match _strip_value_from_type_annotation_that_is_indeed_inside_given_identifier(
+                            annotation, imports_result, identifier_name
+                        ):
+                            case "":
                                 yield node, f"{left} = {right}"
-                            case _, new_annotation:
+                            case str(new_annotation):
                                 yield node, f"{left}: {new_annotation} = {right}"
 
 
+# TODO: make dataclasses frozen and all other stuff  # noqa: FIX002, TD002, TD003
 @dataclass
 class Edit:
     node: SgNode

tests/test_main.py

@@ -779,6 +779,67 @@ def f():
     a = 1
     a = 2
     b: Final = 1
+""",
+            IMPORT_STYLES_TO_IMPORT_CONFIGS["final"],
+        ),
+        (
+            """
+from typing import Final
+
+def f():
+    a: Final[
+        int
+    ] = 1
+---
+from typing import Final
+
+def f():
+    a: Final[
+        int
+    ] = 1
+""",
+            IMPORT_STYLES_TO_IMPORT_CONFIGS["final"],
+        ),
+        (
+            """
+from typing import Annotated
+
+def f():
+    a: Annotated[
+        int,
+        "hello",
+    ] = 1
+    a = 2
+---
+from typing import Annotated
+
+def f():
+    a: Annotated[
+        int,
+        "hello",
+    ] = 1
+    a = 2
+""",
+            IMPORT_STYLES_TO_IMPORT_CONFIGS["final"],
+        ),
+        (
+            """
+from typing import Annotated
+
+def f():
+    a: Annotated[
+        int,
+        "hello"
+    ] = 1
+---
+from typing import Final
+from typing import Annotated
+
+def f():
+    a: Final[Annotated[
+        int,
+        "hello"
+    ]] = 1
 """,
             IMPORT_STYLES_TO_IMPORT_CONFIGS["final"],
         ),