Allow to configure import style from extension (#29)

Author: vrslev

.gitignore

@@ -9,3 +9,4 @@ uv.lock
 .pytest_cache
 .ruff_cache
 testing/*.py
+testing/.vscode

auto_typing_final/lsp.py

@@ -1,6 +1,7 @@
+import uuid
 from collections.abc import Iterable
 from importlib.metadata import version
-from typing import cast
+from typing import cast, get_args
 
 import attr
 import cattrs
@@ -9,10 +10,17 @@
 from pygls import server
 from pygls.workspace import TextDocument
 
-from auto_typing_final.transform import IMPORT_MODES_TO_IMPORT_CONFIGS, AddFinal, Edit, ImportMode, make_replacements
+from auto_typing_final.transform import (
+    IMPORT_STYLES_TO_IMPORT_CONFIGS,
+    AddFinal,
+    Edit,
+    ImportConfig,
+    ImportStyle,
+    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]
+IMPORT_CONFIG: ImportConfig | None = None
 
 
 @attr.define
@@ -45,15 +53,17 @@ def make_text_edit(edit: Edit) -> lsp.TextEdit:
 
 
 def make_diagnostics(source: str) -> Iterable[lsp.Diagnostic]:
+    if not IMPORT_CONFIG:
+        return
     result = make_replacements(root=SgRoot(source, "python").root(), import_config=IMPORT_CONFIG)
 
     for replacement in result.replacements:
         if replacement.operation_type == AddFinal:
-            fix_message = f"{LSP_SERVER.name}: Add typing.Final"
-            diagnostic_message = "Missing typing.Final"
+            fix_message = f"{LSP_SERVER.name}: Add {IMPORT_CONFIG.value}"
+            diagnostic_message = f"Missing {IMPORT_CONFIG.value}"
         else:
-            fix_message = f"{LSP_SERVER.name}: Remove typing.Final"
-            diagnostic_message = "Unexpected typing.Final"
+            fix_message = f"{LSP_SERVER.name}: Remove {IMPORT_CONFIG.value}"
+            diagnostic_message = f"Unexpected {IMPORT_CONFIG.value}"
 
         fix = Fix(message=fix_message, text_edits=[make_text_edit(edit) for edit in replacement.edits])
         if result.import_text:
@@ -74,6 +84,8 @@ def make_diagnostics(source: str) -> Iterable[lsp.Diagnostic]:
 
 
 def make_fixall_text_edits(source: str) -> Iterable[lsp.TextEdit]:
+    if not IMPORT_CONFIG:
+        return
     result = make_replacements(root=SgRoot(source, "python").root(), import_config=IMPORT_CONFIG)
 
     for replacement in result.replacements:
@@ -98,7 +110,35 @@ def make_workspace_edit(text_document: TextDocument, text_edits: list[lsp.TextEd
 
 
 @LSP_SERVER.feature(lsp.INITIALIZE)
-def initialize(params: lsp.InitializeParams) -> None: ...  # noqa: ARG001
+async def initialize(_: lsp.InitializeParams) -> None: ...
+
+
+@LSP_SERVER.feature(lsp.INITIALIZED)
+async def initialized(_: lsp.InitializedParams) -> None:
+    await LSP_SERVER.register_capability_async(
+        params=lsp.RegistrationParams(
+            registrations=[
+                lsp.Registration(
+                    id=str(uuid.uuid4()),
+                    method=lsp.WORKSPACE_DID_CHANGE_CONFIGURATION,
+                    register_options=lsp.DidChangeConfigurationRegistrationOptions(section=LSP_SERVER.name),
+                )
+            ]
+        )
+    )
+
+
+@LSP_SERVER.feature(lsp.WORKSPACE_DID_CHANGE_CONFIGURATION)
+def workspace_did_change_configuration(params: lsp.DidChangeConfigurationParams) -> None:
+    if (
+        isinstance(params.settings, dict)
+        and (settings := params.settings.get(LSP_SERVER.name))
+        and isinstance(settings, dict)
+        and (import_style := settings.get("import-style"))
+        and (import_style in get_args(ImportStyle))
+    ):
+        global IMPORT_CONFIG  # noqa: PLW0603
+        IMPORT_CONFIG = IMPORT_STYLES_TO_IMPORT_CONFIGS[import_style]
 
 
 @LSP_SERVER.feature(lsp.TEXT_DOCUMENT_DID_OPEN)

auto_typing_final/main.py

@@ -3,15 +3,11 @@
 from collections.abc import Iterable
 from difflib import unified_diff
 from pathlib import Path
+from typing import get_args
 
 from ast_grep_py import SgRoot
 
-from auto_typing_final.transform import (
-    IMPORT_MODES_TO_IMPORT_CONFIGS,
-    ImportConfig,
-    ImportMode,
-    make_replacements,
-)
+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:
@@ -49,10 +45,10 @@ def main() -> int:
     parser = argparse.ArgumentParser()
     parser.add_argument("files", type=Path, nargs="*")
     parser.add_argument("--check", action="store_true")
-    parser.add_argument("--import-mode", type=ImportMode, default=ImportMode.typing_final)
+    parser.add_argument("--import-style", type=str, choices=get_args(ImportStyle), default="typing-final")
 
     args = parser.parse_args()
-    import_config = IMPORT_MODES_TO_IMPORT_CONFIGS[args.import_mode]
+    import_config = IMPORT_STYLES_TO_IMPORT_CONFIGS[args.import_style]
 
     has_changes = False
 

auto_typing_final/transform.py

@@ -1,16 +1,13 @@
 import re
 from collections.abc import Iterable
 from dataclasses import dataclass
-from enum import Enum
+from typing import Literal
 
 from ast_grep_py import SgNode
 
 from auto_typing_final.finder import find_all_definitions_in_functions, has_global_identifier_with_name
 
-
-class ImportMode(Enum):
-    typing_final = "typing-final"
-    final = "final"
+ImportStyle = Literal["typing-final", "final"]
 
 
 @dataclass
@@ -21,14 +18,14 @@ class ImportConfig:
     import_identifier: str
 
 
-IMPORT_MODES_TO_IMPORT_CONFIGS = {
-    ImportMode.typing_final: ImportConfig(
+IMPORT_STYLES_TO_IMPORT_CONFIGS: dict[ImportStyle, ImportConfig] = {
+    "typing-final": ImportConfig(
         value="typing.Final",
         outer_regex=re.compile(r"typing\.Final\[(.*)\]{1}"),
         import_text="import typing",
         import_identifier="typing",
     ),
-    ImportMode.final: ImportConfig(
+    "final": ImportConfig(
         value="Final",
         outer_regex=re.compile(r"Final\[(.*)\]{1}"),
         import_text="from typing import Final",

package.json

@@ -24,7 +24,21 @@
 				"category": "auto-typing-final",
 				"command": "auto-typing-final.restart"
 			}
-		]
+		],
+		"configuration": {
+			"properties": {
+				"auto-typing-final.import-style": {
+					"default": "typing-final",
+					"enum": ["typing-final", "final"],
+					"enumDescriptions": [
+						"Use `import typing` and `typing.Final`",
+						"Use `from typing import Final` and `Final`"
+					],
+					"description": "Import style",
+					"scope": "resource"
+				}
+			}
+		}
 	},
 	"scripts": {
 		"lint": "biome format --write .",

tests/test_main.py

@@ -1,14 +1,10 @@
 import pytest
 
 from auto_typing_final.main import transform_file_content
-from auto_typing_final.transform import (
-    IMPORT_MODES_TO_IMPORT_CONFIGS,
-    ImportConfig,
-    ImportMode,
-)
+from auto_typing_final.transform import IMPORT_STYLES_TO_IMPORT_CONFIGS, ImportConfig
 
 
-@pytest.mark.parametrize("import_config", IMPORT_MODES_TO_IMPORT_CONFIGS.values())
+@pytest.mark.parametrize("import_config", IMPORT_STYLES_TO_IMPORT_CONFIGS.values())
 @pytest.mark.parametrize(
     ("before", "after"),
     [
@@ -591,7 +587,7 @@ def foo(self):
     ],
 )
 def test_transform_file_content(case: str) -> None:
-    import_config = IMPORT_MODES_TO_IMPORT_CONFIGS[ImportMode.typing_final]
+    import_config = IMPORT_STYLES_TO_IMPORT_CONFIGS["typing-final"]
     before, _, after = case.partition("---")
     assert (
         transform_file_content(f"{import_config.import_text}\n" + before.strip(), import_config=import_config)
@@ -656,5 +652,5 @@ def f():
 )
 def test_add_import(case: str) -> None:
     before, _, after = case.partition("---")
-    import_config = IMPORT_MODES_TO_IMPORT_CONFIGS[ImportMode.typing_final]
+    import_config = IMPORT_STYLES_TO_IMPORT_CONFIGS["typing-final"]
     assert transform_file_content(before.strip(), import_config=import_config) == after.strip()