Support from typing import Final style (#28)

Author: vrslev

auto_typing_final/finder.py

@@ -28,62 +28,51 @@
 }
 
 
-def last_child_of_type(node: SgNode, type_: str) -> SgNode | None:
+def _get_last_child_of_type(node: SgNode, type_: str) -> SgNode | None:
     return last_child if (children := node.children()) and (last_child := children[-1]).kind() == type_ else None
 
 
-def find_identifiers_in_children(node: SgNode) -> Iterable[SgNode]:
+def _find_identifiers_in_children(node: SgNode) -> Iterable[SgNode]:
     for child in node.children():
         if child.kind() == "identifier":
             yield child
 
 
-def is_inside_inner_function(root: SgNode, node: SgNode) -> bool:
+def _is_inside_inner_function(root: SgNode, node: SgNode) -> bool:
     for ancestor in node.ancestors():
         if ancestor.kind() == "function_definition":
             return ancestor != root
     return False
 
 
-def is_inside_inner_function_or_class(root: SgNode, node: SgNode) -> bool:
+def _is_inside_inner_function_or_class(root: SgNode, node: SgNode) -> bool:
     for ancestor in node.ancestors():
         if ancestor.kind() in {"function_definition", "class_definition"}:
             return ancestor != root
     return False
 
 
-def find_identifiers_in_function_parameter(node: SgNode) -> Iterable[SgNode]:
-    match node.kind():
-        case "default_parameter" | "typed_default_parameter":
-            if name := node.field("name"):
-                yield name
-        case "identifier":
-            yield node
-        case _:
-            yield from find_identifiers_in_children(node)
-
-
-def find_identifiers_in_import(node: SgNode) -> Iterable[SgNode]:
+def _find_identifiers_in_import_statement(node: SgNode) -> Iterable[SgNode]:
     match tuple((child.kind(), child) for child in node.children()):
         case (("from", _), _, ("import", _), *name_nodes) | (("import", _), *name_nodes):
             for kind, name_node in name_nodes:
                 match kind:
                     case "dotted_name":
-                        if identifier := last_child_of_type(name_node, "identifier"):
+                        if identifier := _get_last_child_of_type(name_node, "identifier"):
                             yield identifier
                     case "aliased_import":
                         if alias := name_node.field("alias"):
                             yield alias
 
 
-def find_identifiers_in_function_body(node: SgNode) -> Iterable[SgNode]:  # noqa: C901, PLR0912
+def _find_identifiers_made_by_node(node: SgNode) -> Iterable[SgNode]:  # noqa: C901, PLR0912
     match node.kind():
         case "assignment" | "augmented_assignment":
             if not (left := node.field("left")):
                 return
             match left.kind():
                 case "pattern_list" | "tuple_pattern":
-                    yield from find_identifiers_in_children(left)
+                    yield from _find_identifiers_in_children(left)
                 case "identifier":
                     yield left
         case "named_expression":
@@ -94,18 +83,18 @@ def find_identifiers_in_function_body(node: SgNode) -> Iterable[SgNode]:  # noqa
                 yield name
             for function in node.find_all(kind="function_definition"):
                 for nonlocal_statement in node.find_all(kind="nonlocal_statement"):
-                    if is_inside_inner_function(root=function, node=nonlocal_statement):
+                    if _is_inside_inner_function(root=function, node=nonlocal_statement):
                         continue
-                    yield from find_identifiers_in_children(nonlocal_statement)
+                    yield from _find_identifiers_in_children(nonlocal_statement)
         case "function_definition":
             if name := node.field("name"):
                 yield name
             for nonlocal_statement in node.find_all(kind="nonlocal_statement"):
-                if is_inside_inner_function(root=node, node=nonlocal_statement):
+                if _is_inside_inner_function(root=node, node=nonlocal_statement):
                     continue
-                yield from find_identifiers_in_children(nonlocal_statement)
+                yield from _find_identifiers_in_children(nonlocal_statement)
         case "import_from_statement" | "import_statement":
-            yield from find_identifiers_in_import(node)
+            yield from _find_identifiers_in_import_statement(node)
         case "as_pattern":
             match tuple((child.kind(), child) for child in node.children()):
                 case (
@@ -116,62 +105,66 @@ def find_identifiers_in_function_body(node: SgNode) -> Iterable[SgNode]:  # noqa
         case "keyword_pattern":
             match tuple((child.kind(), child) for child in node.children()):
                 case (("identifier", _), ("=", _), ("dotted_name", alias)):
-                    if identifier := last_child_of_type(alias, "identifier"):
+                    if identifier := _get_last_child_of_type(alias, "identifier"):
                         yield identifier
         case "list_pattern" | "tuple_pattern":
             for child in node.children():
                 if (
                     child.kind() == "case_pattern"
-                    and (dotted_name := last_child_of_type(child, "dotted_name"))
-                    and (identifier := last_child_of_type(dotted_name, "identifier"))
+                    and (dotted_name := _get_last_child_of_type(child, "dotted_name"))
+                    and (identifier := _get_last_child_of_type(dotted_name, "identifier"))
                 ):
                     yield identifier
         case "splat_pattern" | "global_statement" | "nonlocal_statement":
-            yield from find_identifiers_in_children(node)
+            yield from _find_identifiers_in_children(node)
         case "dict_pattern":
             for child in node.children():
                 if (
                     child.kind() == "case_pattern"
                     and (previous_child := child.prev())
                     and previous_child.kind() == ":"
-                    and (dotted_name := last_child_of_type(child, "dotted_name"))
-                    and (identifier := last_child_of_type(dotted_name, "identifier"))
+                    and (dotted_name := _get_last_child_of_type(child, "dotted_name"))
+                    and (identifier := _get_last_child_of_type(dotted_name, "identifier"))
                 ):
                     yield identifier
         case "for_statement":
             if left := node.field("left"):
                 yield from left.find_all(kind="identifier")
 
 
-def find_definitions_in_scope_grouped_by_name(root: SgNode) -> dict[str, list[SgNode]]:
-    definition_map = defaultdict(list)
-
-    if parameters := root.field("parameters"):
-        for parameter in parameters.children():
-            for identifier in find_identifiers_in_function_parameter(parameter):
-                definition_map[identifier.text()].append(parameter)
-
-    for node in root.find_all(DEFINITION_RULE):
-        if is_inside_inner_function_or_class(root, node) or node == root:
+def _find_identifiers_in_scope(node: SgNode) -> Iterable[tuple[SgNode, SgNode]]:
+    for child in node.find_all(DEFINITION_RULE):
+        if _is_inside_inner_function_or_class(node, child) or child == node:
             continue
-        for identifier in find_identifiers_in_function_body(node):
-            definition_map[identifier.text()].append(node)
+        for identifier in _find_identifiers_made_by_node(child):
+            yield identifier, child
+
 
-    return definition_map
+def _find_identifiers_in_function_parameter(node: SgNode) -> Iterable[SgNode]:
+    match node.kind():
+        case "default_parameter" | "typed_default_parameter":
+            if name := node.field("name"):
+                yield name
+        case "identifier":
+            yield node
+        case _:
+            yield from _find_identifiers_in_children(node)
 
 
-def find_definitions_in_module(root: SgNode) -> Iterable[list[SgNode]]:
+def find_all_definitions_in_functions(root: SgNode) -> Iterable[list[SgNode]]:
     for function in root.find_all(kind="function_definition"):
-        yield from find_definitions_in_scope_grouped_by_name(function).values()
+        definition_map = defaultdict(list)
 
+        if parameters := function.field("parameters"):
+            for parameter in parameters.children():
+                for identifier in _find_identifiers_in_function_parameter(parameter):
+                    definition_map[identifier.text()].append(parameter)
 
-def has_global_import_with_name(root: SgNode, name: str) -> bool:
-    for import_statement in root.find_all(
-        {"rule": {"any": [{"kind": "import_from_statement"}, {"kind": "import_statement"}]}}
-    ):
-        if is_inside_inner_function_or_class(root, import_statement):
-            continue
-        for identifier in find_identifiers_in_import(import_statement):
-            if identifier.text() == name:
-                return True
-    return False
+        for identifier, node in _find_identifiers_in_scope(function):
+            definition_map[identifier.text()].append(node)
+
+        yield from definition_map.values()
+
+
+def has_global_identifier_with_name(root: SgNode, name: str) -> bool:
+    return name in {identifier.text() for identifier, _ in _find_identifiers_in_scope(root)}

auto_typing_final/lsp.py

@@ -9,10 +9,10 @@
 from pygls import server
 from pygls.workspace import TextDocument
 
-from auto_typing_final.finder import has_global_import_with_name
-from auto_typing_final.transform import AddFinal, AppliedOperation, make_operations_from_root
+from auto_typing_final.transform import IMPORT_MODES_TO_IMPORT_CONFIGS, AddFinal, Edit, ImportMode, make_replacements
 
 LSP_SERVER = server.LanguageServer(name="auto-typing-final", version=version("auto-typing-final"), max_workers=5)
+IMPORT_CONFIG = IMPORT_MODES_TO_IMPORT_CONFIGS[ImportMode.typing_final]
 
 
 @attr.define
@@ -26,43 +26,40 @@ class DiagnosticData:
     fix: Fix
 
 
-def make_typing_import() -> lsp.TextEdit:
+def make_import_text_edit(import_text: str) -> lsp.TextEdit:
     return lsp.TextEdit(
         range=lsp.Range(start=lsp.Position(line=0, character=0), end=lsp.Position(line=0, character=0)),
-        new_text="import typing\n",
+        new_text=f"{import_text}\n",
     )
 
 
-def make_diagnostic_text_edits(applied_operation: AppliedOperation) -> Iterable[lsp.TextEdit]:
-    for applied_edit in applied_operation.edits:
-        node_range = applied_edit.node.range()
-        yield lsp.TextEdit(
-            range=lsp.Range(
-                start=lsp.Position(line=node_range.start.line, character=node_range.start.column),
-                end=lsp.Position(line=node_range.end.line, character=node_range.end.column),
-            ),
-            new_text=applied_edit.edit.inserted_text,
-        )
+def make_text_edit(edit: Edit) -> lsp.TextEdit:
+    node_range = edit.node.range()
+    return lsp.TextEdit(
+        range=lsp.Range(
+            start=lsp.Position(line=node_range.start.line, character=node_range.start.column),
+            end=lsp.Position(line=node_range.end.line, character=node_range.end.column),
+        ),
+        new_text=edit.new_text,
+    )
 
 
 def make_diagnostics(source: str) -> Iterable[lsp.Diagnostic]:
-    root = SgRoot(source, "python").root()
-    has_import = has_global_import_with_name(root, "typing")
+    result = make_replacements(root=SgRoot(source, "python").root(), import_config=IMPORT_CONFIG)
 
-    for applied_operation in make_operations_from_root(root):
-        if isinstance(applied_operation.operation, AddFinal):
+    for replacement in result.replacements:
+        if replacement.operation_type == AddFinal:
             fix_message = f"{LSP_SERVER.name}: Add typing.Final"
             diagnostic_message = "Missing typing.Final"
         else:
             fix_message = f"{LSP_SERVER.name}: Remove typing.Final"
             diagnostic_message = "Unexpected typing.Final"
 
-        fix = Fix(message=fix_message, text_edits=list(make_diagnostic_text_edits(applied_operation)))
-
-        if isinstance(applied_operation.operation, AddFinal) and not has_import:
-            fix.text_edits.append(make_typing_import())
+        fix = Fix(message=fix_message, text_edits=[make_text_edit(edit) for edit in replacement.edits])
+        if result.import_text:
+            fix.text_edits.append(make_import_text_edit(result.import_text))
 
-        for applied_edit in applied_operation.edits:
+        for applied_edit in replacement.edits:
             node_range = applied_edit.node.range()
             yield lsp.Diagnostic(
                 range=lsp.Range(
@@ -77,36 +74,26 @@ def make_diagnostics(source: str) -> Iterable[lsp.Diagnostic]:
 
 
 def make_fixall_text_edits(source: str) -> Iterable[lsp.TextEdit]:
-    root = SgRoot(source, "python").root()
-    has_import = has_global_import_with_name(root, "typing")
-    has_add_final_operation = False
-
-    for applied_operation in make_operations_from_root(root):
-        if isinstance(applied_operation.operation, AddFinal):
-            has_add_final_operation = True
-        yield from make_diagnostic_text_edits(applied_operation)
-
-    if has_add_final_operation and not has_import:
-        yield make_typing_import()
-
-
-def make_quickfix_action(diagnostic: lsp.Diagnostic, text_document: TextDocument) -> lsp.CodeAction:
-    data = cattrs.structure(diagnostic.data, DiagnosticData)
-    return lsp.CodeAction(
-        title=data.fix.message,
-        kind=lsp.CodeActionKind.QuickFix,
-        data=text_document.uri,
-        edit=lsp.WorkspaceEdit(
-            document_changes=[
-                lsp.TextDocumentEdit(
-                    text_document=lsp.OptionalVersionedTextDocumentIdentifier(
-                        uri=text_document.uri, version=text_document.version
-                    ),
-                    edits=data.fix.text_edits,  # type: ignore[arg-type]
-                )
-            ]
-        ),
-        diagnostics=[diagnostic],
+    result = make_replacements(root=SgRoot(source, "python").root(), import_config=IMPORT_CONFIG)
+
+    for replacement in result.replacements:
+        for edit in replacement.edits:
+            yield make_text_edit(edit)
+
+    if result.import_text:
+        yield make_import_text_edit(result.import_text)
+
+
+def make_workspace_edit(text_document: TextDocument, text_edits: list[lsp.TextEdit]) -> lsp.WorkspaceEdit:
+    return lsp.WorkspaceEdit(
+        document_changes=[
+            lsp.TextDocumentEdit(
+                text_document=lsp.OptionalVersionedTextDocumentIdentifier(
+                    uri=text_document.uri, version=text_document.version
+                ),
+                edits=text_edits,  # type: ignore[arg-type]
+            )
+        ]
     )
 
 
@@ -141,11 +128,31 @@ def code_action(params: lsp.CodeActionParams) -> list[lsp.CodeAction] | None:
 
     if lsp.CodeActionKind.QuickFix in requested_kinds:
         text_document = LSP_SERVER.workspace.get_text_document(params.text_document.uri)
-        actions.extend(
-            make_quickfix_action(diagnostic=diagnostic, text_document=text_document)
-            for diagnostic in params.context.diagnostics
-            if diagnostic.source == LSP_SERVER.name
-        )
+        our_diagnostics = [
+            diagnostic for diagnostic in params.context.diagnostics if diagnostic.source == LSP_SERVER.name
+        ]
+
+        for diagnostic in our_diagnostics:
+            data = cattrs.structure(diagnostic.data, DiagnosticData)
+            actions.append(
+                lsp.CodeAction(
+                    title=data.fix.message,
+                    kind=lsp.CodeActionKind.QuickFix,
+                    edit=make_workspace_edit(text_document=text_document, text_edits=data.fix.text_edits),
+                    diagnostics=[diagnostic],
+                )
+            )
+
+        if our_diagnostics:
+            actions.append(
+                lsp.CodeAction(
+                    title=f"{LSP_SERVER.name}: Fix All",
+                    kind=lsp.CodeActionKind.QuickFix,
+                    data=params.text_document.uri,
+                    edit=None,
+                    diagnostics=params.context.diagnostics,
+                )
+            )
 
     if lsp.CodeActionKind.SourceFixAll in requested_kinds:
         actions.append(
@@ -154,7 +161,7 @@ def code_action(params: lsp.CodeActionParams) -> list[lsp.CodeAction] | None:
                 kind=lsp.CodeActionKind.SourceFixAll,
                 data=params.text_document.uri,
                 edit=None,
-                diagnostics=[],
+                diagnostics=params.context.diagnostics,
             ),
         )
 
@@ -163,9 +170,6 @@ def code_action(params: lsp.CodeActionParams) -> list[lsp.CodeAction] | None:
 
 @LSP_SERVER.feature(lsp.CODE_ACTION_RESOLVE)
 def resolve_code_action(params: lsp.CodeAction) -> lsp.CodeAction:
-    if params.kind != lsp.CodeActionKind.SourceFixAll:
-        return params
-
     text_document = LSP_SERVER.workspace.get_text_document(cast(str, params.data))
     params.edit = lsp.WorkspaceEdit(
         document_changes=[