Refactor to avoid method-level import in hot code

Author: ilevkivskyi

mypy/build.py

@@ -143,7 +143,8 @@
 
 from mypy import errorcodes as codes
 from mypy.config_parser import get_config_module_names, parse_mypy_comments
-from mypy.fixup import node_fixer
+from mypy.fixer_state import fixer_state
+from mypy.fixup import NodeFixer
 from mypy.freetree import free_tree
 from mypy.fscache import FileSystemCache
 from mypy.known_modules import get_known_modules, reset_known_modules_cache
@@ -813,11 +814,9 @@ def __init__(
         self.version_id = version_id
         self.modules: dict[str, MypyFile] = {}
         # Share same modules dictionary with the global fixer state.
-        node_fixer.modules = node_fixer.type_fixer.modules = self.modules
         # We need to set allow_missing when doing a fine-grained cache
         # load because we need to gracefully handle missing modules.
-        allow_missing = self.options.use_fine_grained_cache
-        node_fixer.allow_missing = node_fixer.type_fixer.allow_missing = allow_missing
+        fixer_state.node_fixer = NodeFixer(self.modules, self.options.use_fine_grained_cache)
         self.import_map: dict[str, set[str]] = {}
         self.missing_modules: dict[str, int] = {}
         self.fg_deps_meta: dict[str, FgDepMeta] = {}
@@ -2820,8 +2819,9 @@ def load_tree(self, temporary: bool = False) -> None:
     def fix_cross_refs(self) -> None:
         assert self.tree is not None, "Internal error: method must be called on parsed file only"
         # Do initial lightweight pass fixing TypeInfos and module cross-references.
-        node_fixer.visit_symbol_table(self.tree.names)
-        type_fixer = node_fixer.type_fixer
+        assert fixer_state.node_fixer is not None
+        fixer_state.node_fixer.visit_symbol_table(self.tree.names)
+        type_fixer = fixer_state.node_fixer.type_fixer
         if instance_cache.str_type is not None:
             instance_cache.str_type.accept(type_fixer)
         if instance_cache.function_type is not None:

mypy/fixer_state.py

@@ -0,0 +1,18 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Final
+
+if TYPE_CHECKING:
+    from mypy.fixup import NodeFixer
+
+# This is global mutable state. Don't add anything here unless there's a very
+# good reason. This exists as a separate file to avoid method-level import in
+# hot code in SymbolTableNode.node().
+
+
+class FixerState:
+    def __init__(self) -> None:
+        self.node_fixer: NodeFixer | None = None
+
+
+fixer_state: Final = FixerState()

mypy/fixup.py

@@ -49,13 +49,13 @@
 class NodeFixer(NodeVisitor[None]):
     current_info: TypeInfo | None = None
 
-    def __init__(self) -> None:
-        self.modules: dict[str, MypyFile] = {}
+    def __init__(self, modules: dict[str, MypyFile], allow_missing: bool) -> None:
+        self.modules = modules
         # N.B: we do an allow_missing fixup when fixing up a fine-grained
         # incremental cache load (since there may be cross-refs into deleted
         # modules)
-        self.allow_missing = False
-        self.type_fixer = TypeFixer()
+        self.allow_missing = allow_missing
+        self.type_fixer = TypeFixer(modules, allow_missing)
 
     # NOTE: This method isn't (yet) part of the NodeVisitor API.
     def visit_type_info(self, info: TypeInfo) -> None:
@@ -225,9 +225,9 @@ def visit_type_alias(self, a: TypeAlias) -> None:
 
 
 class TypeFixer(TypeVisitor[None]):
-    def __init__(self) -> None:
-        self.modules: dict[str, MypyFile] = {}
-        self.allow_missing = False
+    def __init__(self, modules: dict[str, MypyFile], allow_missing: bool) -> None:
+        self.modules = modules
+        self.allow_missing = allow_missing
 
     def visit_instance(self, inst: Instance) -> None:
         type_ref = inst.type_ref
@@ -434,6 +434,3 @@ def missing_info(modules: dict[str, MypyFile]) -> TypeInfo:
 def missing_alias() -> TypeAlias:
     suggestion = _SUGGESTION.format("alias")
     return TypeAlias(AnyType(TypeOfAny.special_form), suggestion, "<missing>", line=-1, column=-1)
-
-
-node_fixer: Final = NodeFixer()

mypy/nodes.py

@@ -67,6 +67,7 @@
     write_str_opt_list,
     write_tag,
 )
+from mypy.fixer_state import fixer_state
 from mypy.options import Options
 from mypy.util import is_sunder, is_typeshed_file, short_type
 from mypy.visitor import ExpressionVisitor, NodeVisitor, StatementVisitor
@@ -4820,10 +4821,9 @@ def type(self) -> mypy.types.Type | None:
 
     @property
     def node(self) -> SymbolNode | None:
-        # Late import because of a circular dependency.
-        from mypy.fixup import node_fixer
-
         if self.unfixed:
+            node_fixer = fixer_state.node_fixer
+            assert node_fixer is not None
             if self.cross_ref is not None:
                 node_fixer.resolve_cross_ref(self)
             else:

mypy/plugins/common.py

@@ -3,7 +3,7 @@
 from typing import NamedTuple
 
 from mypy.argmap import map_actuals_to_formals
-from mypy.fixup import node_fixer
+from mypy.fixer_state import fixer_state
 from mypy.nodes import (
     ARG_POS,
     MDEF,
@@ -435,5 +435,6 @@ def add_attribute_to_class(
 
 def deserialize_and_fixup_type(data: str | JsonDict, api: SemanticAnalyzerPluginInterface) -> Type:
     typ = deserialize_type(data)
-    typ.accept(node_fixer.type_fixer)
+    assert fixer_state.node_fixer is not None
+    typ.accept(fixer_state.node_fixer.type_fixer)
     return typ