Ignore global vars (#124, issue: #96)

Author: vrslev

README.md

@@ -15,10 +15,10 @@ Auto-fixer for Python code that adds `typing.Final` annotation to variable assig
 Basically, this, but handles different operations (like usage of `nonlocal`, augmented assignments: `+=`, etc) as well.
 
 - Keeps mypy happy.
-- Ignores global variables to avoid confusion with the type aliases like `Fruit = Apple | Banana`.
-- Ignores class variables: it is common to use `typing.ClassVar` instead of `typing.Final`.
 - Adds global import if it's not imported yet (`import typing`/`from typing import Final`).
 - Inspects one file at a time.
+- Is careful with global variables: adds Final only for uppercase variables, ignores variable that are referenced in `global` statement inside functions in current file, and avoids removing Final when it already was set.
+- Ignores class variables: it is common to use `typing.ClassVar` instead of `typing.Final`.
 
 ## How To Use
 
@@ -34,7 +34,6 @@ or:
 pipx run auto-typing-final .
 ```
 
-
 ### Options
 
 You can specify `--check` flag to check the files instead of actually fixing them:
@@ -52,8 +51,17 @@ auto-typing-final . --import-style typing-final
 - `typing-final` enforces `import typing` and `typing.Final`,
 - `final` enforces `from typing import Final` and `Final`.
 
+Also, you can set `--ignore-global-vars` flag to ignore global variables:
+
+```sh
+auto-typing-final . --ignore-global-vars
+```
+
+### Ignore comment
+
+You can ignore variables by adding `# auto-typing-final: ignore` comment to the line.
 
-## VS Code Extension
+### VS Code Extension
 
 <img width="768" alt="image" src="https://github.com/community-of-python/auto-typing-final/assets/75225148/f1541056-06f5-4caa-8c94-0a5eaf98ba15">
 
@@ -67,7 +75,8 @@ After that, install the extension: https://marketplace.visualstudio.com/items?it
 
 ### Settings
 
-Import style can be configured in settings: `"auto-typing-final.import-style": "typing-final"` or `"auto-typing-final.import-style": "final"`.
+- Import style can be configured in settings: `"auto-typing-final.import-style": "typing-final"` or `"auto-typing-final.import-style": "final"`.
+- Ignore global variables can be configured in settings: `"auto-typing-final.ignore-global-vars": true`.
 
 ### Notes
 

auto_typing_final/finder.py

@@ -6,7 +6,7 @@
 from ast_grep_py import Config, SgNode
 
 # https://github.com/tree-sitter/tree-sitter-python/blob/71778c2a472ed00a64abf4219544edbf8e4b86d7/grammar.js
-DEFINITION_RULE: Config = {
+DEFINITION_RULE: Final[Config] = {
     "rule": {
         "any": [
             {"kind": "assignment"},
@@ -28,6 +28,14 @@
         ]
     }
 }
+GLOBAL_STATEMENTS_RULE: Final[Config] = {
+    "rule": {
+        "any": [
+            {"kind": "global_statement"},
+        ]
+    }
+}
+IGNORE_COMMENT_TEXT: Final = "# auto-typing-final: ignore"
 
 
 def _get_last_child_of_type(node: SgNode, type_: str) -> SgNode | None:
@@ -134,9 +142,20 @@ def _find_identifiers_made_by_node(node: SgNode) -> Iterable[SgNode]:  # noqa: C
                 yield from left.find_all(kind="identifier")
 
 
+def _line_has_ignore_comment(node: SgNode) -> bool:
+    if not (parent := node.parent()):
+        return False
+    if not (grand_parent := parent.parent()):
+        return False
+    for one_child in grand_parent.children():
+        if one_child.kind() == "comment" and IGNORE_COMMENT_TEXT in one_child.text():
+            return True
+    return False
+
+
 def _find_identifiers_in_current_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:
+        if child == node or _is_inside_inner_function_or_class(node, child) or _line_has_ignore_comment(child):
             continue
         for identifier in _find_identifiers_made_by_node(child):
             yield identifier, child
@@ -172,6 +191,19 @@ def has_global_identifier_with_name(root: SgNode, name: str) -> bool:
     return name in {identifier.text() for identifier, _ in _find_identifiers_in_current_scope(root)}
 
 
+def find_global_definitions(root: SgNode) -> Iterable[list[SgNode]]:
+    definitions_by_name: Final = defaultdict(list)
+
+    for identifier, definition_node in _find_identifiers_in_current_scope(root):
+        definitions_by_name[identifier.text()].append(definition_node)
+
+    for one_node in root.find_all(GLOBAL_STATEMENTS_RULE):
+        for one_identifier in _find_identifiers_in_children(one_node):
+            definitions_by_name[one_identifier.text()].append(one_node)
+
+    return definitions_by_name.values()
+
+
 @dataclass(slots=True, kw_only=True)
 class ImportsResult:
     module_aliases: set[str]

auto_typing_final/lsp.py

@@ -47,7 +47,7 @@ def path_from_uri(uri: str) -> Path | None:
     return path_
 
 
-ClientSettings = TypedDict("ClientSettings", {"import-style": ImportStyle})
+ClientSettings = TypedDict("ClientSettings", {"import-style": ImportStyle, "ignore-global-vars": bool})
 FullClientSettings = TypedDict("FullClientSettings", {"auto-typing-final": ClientSettings})
 
 
@@ -80,6 +80,7 @@ class Service:
     ls_name: str
     ignored_paths: list[Path]
     import_config: ImportConfig
+    ignore_global_vars: bool
 
     @staticmethod
     def try_from_settings(ls_name: str, settings: Any) -> "Service | None":  # noqa: ANN401
@@ -93,11 +94,14 @@ def try_from_settings(ls_name: str, settings: Any) -> "Service | None":  # noqa:
             ls_name=ls_name,
             ignored_paths=[executable_path.parent.parent] if executable_path.parent.name == "bin" else [],
             import_config=IMPORT_STYLES_TO_IMPORT_CONFIGS[validated_settings["auto-typing-final"]["import-style"]],
+            ignore_global_vars=validated_settings["auto-typing-final"]["ignore-global-vars"],
         )
 
     def make_diagnostics(self, source: str) -> list[lsp.Diagnostic]:
         replacement_result: Final = make_replacements(
-            root=SgRoot(source, "python").root(), import_config=self.import_config
+            root=SgRoot(source, "python").root(),
+            import_config=self.import_config,
+            ignore_global_vars=self.ignore_global_vars,
         )
         result: Final = []
 
@@ -131,7 +135,9 @@ def make_diagnostics(self, source: str) -> list[lsp.Diagnostic]:
 
     def make_fix_all_text_edits(self, source: str) -> list[lsp.TextEdit | lsp.AnnotatedTextEdit]:
         replacement_result: Final = make_replacements(
-            root=SgRoot(source, "python").root(), import_config=self.import_config
+            root=SgRoot(source, "python").root(),
+            import_config=self.import_config,
+            ignore_global_vars=self.ignore_global_vars,
         )
         result: Final[list[lsp.TextEdit | lsp.AnnotatedTextEdit]] = [
             make_text_edit(edit) for replacement in replacement_result.replacements for edit in replacement.edits
@@ -150,7 +156,7 @@ class CustomLanguageServer(LanguageServer):
     service: Service | None = None
 
 
-LSP_SERVER = CustomLanguageServer(name="auto-typing-final", version=version("auto-typing-final"), max_workers=5)
+LSP_SERVER: Final = CustomLanguageServer(name="auto-typing-final", version=version("auto-typing-final"), max_workers=5)
 
 
 @LSP_SERVER.feature(lsp.INITIALIZE)

auto_typing_final/main.py

@@ -10,9 +10,9 @@
 from auto_typing_final.transform import IMPORT_STYLES_TO_IMPORT_CONFIGS, ImportConfig, ImportStyle, make_replacements
 
 
-def transform_file_content(source: str, import_config: ImportConfig) -> str:
+def transform_file_content(source: str, import_config: ImportConfig, ignore_global_vars: bool) -> str:  # noqa: FBT001
     root: Final = SgRoot(source, "python").root()
-    result: Final = make_replacements(root, import_config)
+    result: Final = make_replacements(root, import_config, ignore_global_vars)
     new_text: Final = root.commit_edits(
         [edit.node.replace(edit.new_text) for replacement in result.replacements for edit in replacement.edits]
     )
@@ -46,6 +46,9 @@ def main() -> int:
     parser.add_argument("files", type=Path, nargs="*", default=[Path()])
     parser.add_argument("--check", action="store_true")
     parser.add_argument("--import-style", type=str, choices=get_args(ImportStyle), default="typing-final")
+    parser.add_argument(
+        "--ignore-global-vars", action="store_true", help="Ignore global variables when applying Final annotations"
+    )
 
     args: Final = parser.parse_args()
     import_config: Final = IMPORT_STYLES_TO_IMPORT_CONFIGS[args.import_style]
@@ -55,7 +58,9 @@ def main() -> int:
     for path in find_all_source_files(args.files):
         with path.open(open_mode) as file:
             source = file.read()
-            transformed_content = transform_file_content(source=source, import_config=import_config)
+            transformed_content = transform_file_content(
+                source=source, import_config=import_config, ignore_global_vars=args.ignore_global_vars
+            )
             if source == transformed_content:
                 continue
             changed_files_count += 1

auto_typing_final/transform.py

@@ -1,3 +1,4 @@
+import typing
 from collections.abc import Iterable
 from dataclasses import dataclass
 from typing import Final, Literal
@@ -7,6 +8,7 @@
 from auto_typing_final.finder import (
     ImportsResult,
     find_all_definitions_in_functions,
+    find_global_definitions,
     find_imports_of_identifier_in_scope,
     has_global_identifier_with_name,
 )
@@ -20,10 +22,11 @@ class ImportConfig:
 
 
 ImportStyle = Literal["typing-final", "final"]
-IMPORT_STYLES_TO_IMPORT_CONFIGS: dict[ImportStyle, ImportConfig] = {
+IMPORT_STYLES_TO_IMPORT_CONFIGS: Final[dict[ImportStyle, ImportConfig]] = {
     "typing-final": ImportConfig(value="typing.Final", import_text="import typing", import_identifier="typing"),
     "final": ImportConfig(value="Final", import_text="from typing import Final", import_identifier="Final"),
 }
+IGNORED_DEFINITION_PATTERNS: typing.Final = {"TypeVar", "ParamSpec"}
 
 
 @dataclass(frozen=True, slots=True)
@@ -87,18 +90,38 @@ def _make_definition_from_definition_node(node: SgNode) -> Definition:
             return OtherDefinition(node)
 
 
-def _make_operation_from_definitions_of_one_name(nodes: list[SgNode]) -> Operation:
+def _should_skip_global_variable(definition: Definition) -> bool:
+    return isinstance(definition, EditableAssignmentWithoutAnnotation | EditableAssignmentWithAnnotation) and (
+        not (definition.left.isupper() and len(definition.left) > 1)
+        or any(one_pattern in definition.right for one_pattern in IGNORED_DEFINITION_PATTERNS)
+    )
+
+
+def _make_operation_from_definitions_of_one_name(nodes: list[SgNode], ignore_global_vars: bool) -> Operation | None:  # noqa: FBT001
     value_definitions: Final[list[Definition]] = []
     has_node_inside_loop = False
+    has_global_scope_definition = False
 
     for node in nodes:
         if any(ancestor.kind() in {"for_statement", "while_statement"} for ancestor in node.ancestors()):
             has_node_inside_loop = True
+
+        if all(ancestor.kind() != "function_definition" for ancestor in node.ancestors()):
+            has_global_scope_definition = True
+
         value_definitions.append(_make_definition_from_definition_node(node))
 
     if has_node_inside_loop:
         return RemoveFinal(value_definitions)
 
+    if (
+        not ignore_global_vars
+        and has_global_scope_definition
+        and value_definitions
+        and _should_skip_global_variable(value_definitions[0])
+    ):
+        return None
+
     match value_definitions:
         case [definition]:
             return AddFinal(definition)
@@ -192,13 +215,14 @@ class MakeReplacementsResult:
     import_text: str | None
 
 
-def make_replacements(root: SgNode, import_config: ImportConfig) -> MakeReplacementsResult:
+def make_replacements(root: SgNode, import_config: ImportConfig, ignore_global_vars: bool) -> MakeReplacementsResult:  # noqa: FBT001
     replacements: Final = []
     has_added_final = False
     imports_result: Final = find_imports_of_identifier_in_scope(root, module_name="typing", identifier_name="Final")
 
     for current_definitions in find_all_definitions_in_functions(root):
-        operation = _make_operation_from_definitions_of_one_name(current_definitions)
+        if not (operation := _make_operation_from_definitions_of_one_name(current_definitions, ignore_global_vars)):
+            continue
         edits = [
             Edit(node=node, new_text=new_text)
             for node, new_text in _make_changed_text_from_operation(
@@ -215,6 +239,29 @@ def make_replacements(root: SgNode, import_config: ImportConfig) -> MakeReplacem
 
         replacements.append(Replacement(operation_type=operation_type, edits=edits))
 
+    if not ignore_global_vars:
+        for current_definitions in find_global_definitions(root):
+            if (
+                not (operation := _make_operation_from_definitions_of_one_name(current_definitions, ignore_global_vars))
+                or (operation_type := type(operation)) == RemoveFinal
+            ):
+                continue
+
+            edits = [
+                Edit(node=node, new_text=new_text)
+                for node, new_text in _make_changed_text_from_operation(
+                    operation=operation,
+                    final_value=import_config.value,
+                    imports_result=imports_result,
+                    identifier_name="Final",
+                )
+                if node.text() != new_text
+            ]
+            if not edits:
+                continue
+            has_added_final = True
+            replacements.append(Replacement(operation_type=operation_type, edits=edits))
+
     return MakeReplacementsResult(
         replacements=replacements,
         import_text=(

package.json

@@ -36,6 +36,12 @@
 					],
 					"description": "Import style",
 					"scope": "resource"
+				},
+				"auto-typing-final.ignore-global-vars": {
+					"default": false,
+					"type": "boolean",
+					"description": "Do not add Final to global variables.",
+					"scope": "resource"
 				}
 			}
 		}

tests/test_ignore_comment.py

@@ -0,0 +1,28 @@
+from typing import Final
+
+import pytest
+
+from auto_typing_final.main import transform_file_content
+from auto_typing_final.transform import IMPORT_STYLES_TO_IMPORT_CONFIGS
+
+
+@pytest.mark.parametrize(
+    ("before", "after"),
+    [
+        ("VAR_WITH_COMMENT = 1 # some comment", "VAR_WITH_COMMENT: Final = 1 # some comment"),
+        ("IGNORED_VAR = 1 # auto-typing-final: ignore", "IGNORED_VAR = 1 # auto-typing-final: ignore"),
+        ("IGNORED_VAR: Final = 1 # auto-typing-final: ignore", "IGNORED_VAR: Final = 1 # auto-typing-final: ignore"),
+        (
+            "IGNORED_VAR: Final = 1 # auto-typing-final: ignore  # some comment",
+            "IGNORED_VAR: Final = 1 # auto-typing-final: ignore  # some comment",
+        ),
+        ("def foo():\n a = 1", "def foo():\n a: Final = 1"),
+        ("def foo():\n a = 1  # auto-typing-final: ignore", "def foo():\n a = 1  # auto-typing-final: ignore"),
+    ],
+)
+def test_ignore_comment(before: str, after: str) -> None:
+    import_config: Final = IMPORT_STYLES_TO_IMPORT_CONFIGS["final"]
+    result: Final = transform_file_content(
+        f"{import_config.import_text}\n" + before.strip(), import_config=import_config, ignore_global_vars=False
+    )
+    assert result == f"{import_config.import_text}\n" + after.strip()