Merge pull request #331 from reagento/feature/exec-type-checking

add exec_type_checking

docs/changelog/fragments/288.feature.rst

@@ -0,0 +1,2 @@
+Add public api for :func:`.type_tools.exec_type_checking`
+to deal with cyclic references by executing ``if TYPE_CHECKING:`` constructs

src/adaptix/_internal/type_tools/__init__.py

@@ -23,3 +23,4 @@
     make_norm_type,
     normalize_type,
 )
+from .type_evaler import exec_type_checking, make_fragments_collector

src/adaptix/_internal/type_tools/type_evaler.py

@@ -0,0 +1,55 @@
+import ast
+import inspect
+from collections.abc import Callable, Sequence
+from types import ModuleType
+
+
+def make_fragments_collector(*, typing_modules: Sequence[str]) -> Callable[[ast.Module], list[ast.stmt]]:
+    def check_condition(expr: ast.expr) -> bool:
+        # searches for `TYPE_CHECKING`
+        if (
+            isinstance(expr, ast.Name)
+            and isinstance(expr.ctx, ast.Load)
+            and expr.id == "TYPE_CHECKING"
+        ):
+            return True
+
+        # searches for `typing.TYPE_CHECKING`
+        if (  # noqa: SIM103
+            isinstance(expr, ast.Attribute)
+            and expr.attr == "TYPE_CHECKING"
+            and isinstance(expr.ctx, ast.Load)
+            and isinstance(expr.value, ast.Name)
+            and expr.value.id in typing_modules
+            and isinstance(expr.value.ctx, ast.Load)
+        ):
+            return True
+        return False
+
+    def collect_type_checking_only_fragments(module: ast.Module) -> list[ast.stmt]:
+        fragments = []
+        for stmt in module.body:
+            if isinstance(stmt, ast.If) and not stmt.orelse and check_condition(stmt.test):
+                fragments.extend(stmt.body)
+
+        return fragments
+
+    return collect_type_checking_only_fragments
+
+
+default_collector = make_fragments_collector(typing_modules=["typing"])
+
+
+def exec_type_checking(
+    module: ModuleType,
+    *,
+    collector: Callable[[ast.Module], list[ast.stmt]] = default_collector,
+) -> None:
+    source = inspect.getsource(module)
+    fragments = collector(ast.parse(source))
+    code = compile(ast.Module(fragments, type_ignores=[]), f"<exec_type_checking of {module}>", "exec")
+    namespace = module.__dict__.copy()
+    exec(code, namespace)  # noqa: S102
+    for k, v in namespace.items():
+        if not hasattr(module, k):
+            setattr(module, k, v)

src/adaptix/type_tools/__init__.py

@@ -0,0 +1,6 @@
+from adaptix._internal.type_tools import exec_type_checking, make_fragments_collector
+
+__all__ = (
+    "exec_type_checking",
+    "make_fragments_collector",
+)

tests/tests_helpers/tests_helpers/misc.py

@@ -1,4 +1,5 @@
 import dataclasses
+import importlib.util
 import inspect
 import re
 import runpy
@@ -189,6 +190,15 @@ def load_namespace(
     return SimpleNamespace(**ns_dict)
 
 
+@contextmanager
+def temp_module(module: ModuleType):
+    sys.modules[module.__name__] = module
+    try:
+        yield
+    finally:
+        sys.modules.pop(module.__name__, None)
+
+
 @contextmanager
 def load_namespace_keeping_module(
     file_name: str,
@@ -202,11 +212,19 @@ def load_namespace_keeping_module(
     module = ModuleType(run_name)
     for attr, value in ns.__dict__.items():
         setattr(module, attr, value)
-    sys.modules[run_name] = module
-    try:
+
+    with temp_module(module):
         yield ns
-    finally:
-        sys.modules.pop(run_name, None)
+
+
+def import_local_module(file_path: Path, name: Optional[str] = None) -> ModuleType:
+    if name is None:
+        name = file_path.stem
+
+    spec = importlib.util.spec_from_file_location(name, file_path)
+    module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(module)
+    return module
 
 
 def with_notes(exc: E, *notes: Union[str, list[str]]) -> E:

tests/unit/type_tools/data_type_checking.py

@@ -0,0 +1,16 @@
+import typing
+from collections.abc import Sequence
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    IntSeq = Sequence[int]
+
+
+if typing.TYPE_CHECKING:
+    StrSeq = Sequence[str]
+
+
+class Foo:
+    a: bool
+    b: "IntSeq"
+    c: "StrSeq"

tests/unit/type_tools/test_type_evaler.py

@@ -0,0 +1,18 @@
+from collections.abc import Sequence
+from pathlib import Path
+
+from tests_helpers.misc import import_local_module, temp_module
+
+from adaptix._internal.type_tools import get_all_type_hints
+from adaptix.type_tools import exec_type_checking
+
+
+def test_exec_type_checking():
+    module = import_local_module(Path(__file__).with_name("data_type_checking.py"))
+    with temp_module(module):
+        exec_type_checking(module)
+        assert get_all_type_hints(module.Foo) == {
+            "a": bool,
+            "b": Sequence[int],
+            "c": Sequence[str],
+        }