Some refactor

python · ilevkivskyi · Mar 17, 2026 · Mar 15, 2026 · Mar 15, 2026 · Mar 15, 2026
commit 5559cb7e1d436f95e1091ad4235c0660c4bad526
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -130,6 +130,7 @@
     RefExpr,
     ReturnStmt,
     SetExpr,
+    SplittingVisitor,
     StarExpr,
     Statement,
     StrExpr,
@@ -319,7 +320,7 @@ def __exit__(self, exc_type: object, exc_val: object, exc_tb: object) -> Literal
         return False
 
 
-class TypeChecker(NodeVisitor[None], TypeCheckerSharedApi):
+class TypeChecker(NodeVisitor[None], TypeCheckerSharedApi, SplittingVisitor):
     """Mypy type checker.
 
     Type check mypy source files that have been semantically analyzed.
@@ -454,9 +455,6 @@ def __init__(
             or self.path in self.msg.errors.ignored_files
             or (self.options.test_env and self.is_typeshed_stub)
         )
-
-        # If True, process function definitions. If False, don't. This is used
-        # for processing module top levels in fine-grained incremental mode.
         self.recurse_into_functions = True
         # This internal flag is used to track whether we a currently type-checking
         # a final declaration (assignment), so that some errors should be suppressed.
@@ -719,19 +717,6 @@ def accept_loop(
     # Definitions
     #
 
-    @contextmanager
-    def set_recurse_into_functions(self) -> Iterator[None]:
-        """Temporarily set recurse_into_functions to True.
-
-        This is used to process top-level functions/methods as a whole.
-        """
-        old_recurse_into_functions = self.recurse_into_functions
-        self.recurse_into_functions = True
-        try:
-            yield
-        finally:
-            self.recurse_into_functions = old_recurse_into_functions
-
     def visit_overloaded_func_def(self, defn: OverloadedFuncDef) -> None:
         # If a function/method can infer variable types, it should be processed as part
         # of the module top level (i.e. module interface).

diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -6,6 +6,7 @@
 from abc import abstractmethod
 from collections import defaultdict
 from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
 from enum import Enum, unique
 from typing import (
     TYPE_CHECKING,
@@ -5086,6 +5087,26 @@ def read(cls, data: ReadBuffer) -> DataclassTransformSpec:
         return ret
 
 
+@trait
+class SplittingVisitor:
+    # If True, process function definitions. If False, don't. This is used
+    # for processing module top levels in fine-grained incremental mode.
+    recurse_into_functions: bool
+
+    @contextmanager
+    def set_recurse_into_functions(self) -> Iterator[None]:
+        """Temporarily set recurse_into_functions to True.
+
+        This is used to process top-level functions/methods as a whole.
+        """
+        old_recurse_into_functions = self.recurse_into_functions
+        self.recurse_into_functions = True
+        try:
+            yield
+        finally:
+            self.recurse_into_functions = old_recurse_into_functions
+
+
 def get_flags(node: Node, names: list[str]) -> list[str]:
     return [name for name in names if getattr(node, name)]
 

diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -159,6 +159,7 @@
     SetComprehension,
     SetExpr,
     SliceExpr,
+    SplittingVisitor,
     StarExpr,
     Statement,
     StrExpr,
@@ -372,7 +373,7 @@
 
 
 class SemanticAnalyzer(
-    NodeVisitor[None], SemanticAnalyzerInterface, SemanticAnalyzerPluginInterface
+    NodeVisitor[None], SemanticAnalyzerInterface, SemanticAnalyzerPluginInterface, SplittingVisitor
 ):
     """Semantically analyze parsed mypy files.
 
@@ -497,8 +498,6 @@ def __init__(
         self.incomplete_namespaces = incomplete_namespaces
         self.all_exports: list[str] = []
         self.plugin = plugin
-        # If True, process function definitions. If False, don't. This is used
-        # for processing module top levels in fine-grained incremental mode.
         self.recurse_into_functions = True
         self.scope = Scope()
 
@@ -593,19 +592,6 @@ def inside_except_star_block_set(
             if not entering_loop:
                 self.return_stmt_inside_except_star_block = old_return_stmt_flag
 
-    @contextmanager
-    def set_recurse_into_functions(self) -> Iterator[None]:
-        """Temporarily set recurse_into_functions to True.
-
-        This is used to process top-level functions/methods as a whole.
-        """
-        old_recurse_into_functions = self.recurse_into_functions
-        self.recurse_into_functions = True
-        try:
-            yield
-        finally:
-            self.recurse_into_functions = old_recurse_into_functions
-
     #
     # Preparing module (performed before semantic analysis)
     #

diff --git a/mypy/server/aststrip.py b/mypy/server/aststrip.py
@@ -56,6 +56,7 @@
     OpExpr,
     OverloadedFuncDef,
     RefExpr,
+    SplittingVisitor,
     StarExpr,
     SuperExpr,
     TupleExpr,
@@ -79,29 +80,14 @@ def strip_target(node: MypyFile | FuncDef | OverloadedFuncDef) -> None:
         node.accept(visitor)
 
 
-class NodeStripVisitor(TraverserVisitor):
+class NodeStripVisitor(TraverserVisitor, SplittingVisitor):
     def __init__(self) -> None:
         # The current active class.
         self.type: TypeInfo | None = None
         # This is True at class scope, but not in methods.
         self.is_class_body = False
-        # By default, process function definitions. If False, don't -- this is used for
-        # processing module top levels.
         self.recurse_into_functions = True
 
-    @contextmanager
-    def set_recurse_into_functions(self) -> Iterator[None]:
-        """Temporarily set recurse_into_functions to True.
-
-        This is used to process top-level functions/methods as a whole.
-        """
-        old_recurse_into_functions = self.recurse_into_functions
-        self.recurse_into_functions = True
-        try:
-            yield
-        finally:
-            self.recurse_into_functions = old_recurse_into_functions
-
     def strip_file_top_level(self, file_node: MypyFile) -> None:
         """Strip a module top-level (don't recursive into functions)."""
         self.recurse_into_functions = False