Cache information about whether file is typeshed file (#16351)

JukkaL · web-flow · commit c4ab46e6acde · 2023-10-28T15:00:13.000+01:00
We used to check if a file is in typeshed a lot. This seems to speed up
self-check by about 2%, and this should also speed up tests a bit.
diff --git a/mypy/checker.py b/mypy/checker.py
@@ -222,7 +222,7 @@
 from mypy.types_utils import is_overlapping_none, remove_optional, store_argument_type, strip_type
 from mypy.typetraverser import TypeTraverserVisitor
 from mypy.typevars import fill_typevars, fill_typevars_with_any, has_no_typevars
-from mypy.util import is_dunder, is_sunder, is_typeshed_file
+from mypy.util import is_dunder, is_sunder
 from mypy.visitor import NodeVisitor
 
 T = TypeVar("T")
@@ -400,7 +400,7 @@ def __init__(
         self.pass_num = 0
         self.current_node_deferred = False
         self.is_stub = tree.is_stub
-        self.is_typeshed_stub = is_typeshed_file(options.abs_custom_typeshed_dir, path)
+        self.is_typeshed_stub = tree.is_typeshed_file(options)
         self.inferred_attribute_types = None
 
         # If True, process function definitions. If False, don't. This is used
diff --git a/mypy/nodes.py b/mypy/nodes.py
@@ -27,7 +27,7 @@
 
 import mypy.strconv
 from mypy.options import Options
-from mypy.util import short_type
+from mypy.util import is_typeshed_file, short_type
 from mypy.visitor import ExpressionVisitor, NodeVisitor, StatementVisitor
 
 if TYPE_CHECKING:
@@ -283,6 +283,7 @@ class MypyFile(SymbolNode):
         "is_partial_stub_package",
         "plugin_deps",
         "future_import_flags",
+        "_is_typeshed_file",
     )
 
     __match_args__ = ("name", "path", "defs")
@@ -319,6 +320,7 @@ class MypyFile(SymbolNode):
     plugin_deps: dict[str, set[str]]
     # Future imports defined in this file. Populated during semantic analysis.
     future_import_flags: set[str]
+    _is_typeshed_file: bool | None
 
     def __init__(
         self,
@@ -346,6 +348,7 @@ def __init__(
         self.is_cache_skeleton = False
         self.is_partial_stub_package = False
         self.future_import_flags = set()
+        self._is_typeshed_file = None
 
     def local_definitions(self) -> Iterator[Definition]:
         """Return all definitions within the module (including nested).
@@ -371,6 +374,12 @@ def is_package_init_file(self) -> bool:
     def is_future_flag_set(self, flag: str) -> bool:
         return flag in self.future_import_flags
 
+    def is_typeshed_file(self, options: Options) -> bool:
+        # Cache result since this is called a lot
+        if self._is_typeshed_file is None:
+            self._is_typeshed_file = is_typeshed_file(options.abs_custom_typeshed_dir, self.path)
+        return self._is_typeshed_file
+
     def serialize(self) -> JsonDict:
         return {
             ".class": "MypyFile",
diff --git a/mypy/semanal.py b/mypy/semanal.py
@@ -283,14 +283,7 @@
 )
 from mypy.types_utils import is_invalid_recursive_alias, store_argument_type
 from mypy.typevars import fill_typevars
-from mypy.util import (
-    correct_relative_import,
-    is_dunder,
-    is_typeshed_file,
-    module_prefix,
-    unmangle,
-    unnamed_function,
-)
+from mypy.util import correct_relative_import, is_dunder, module_prefix, unmangle, unnamed_function
 from mypy.visitor import NodeVisitor
 
 T = TypeVar("T")
@@ -777,9 +770,7 @@ def file_context(
         self.cur_mod_id = file_node.fullname
         with scope.module_scope(self.cur_mod_id):
             self._is_stub_file = file_node.path.lower().endswith(".pyi")
-            self._is_typeshed_stub_file = is_typeshed_file(
-                options.abs_custom_typeshed_dir, file_node.path
-            )
+            self._is_typeshed_stub_file = file_node.is_typeshed_file(options)
             self.globals = file_node.names
             self.tvar_scope = TypeVarLikeScope()
 
diff --git a/mypy/semanal_main.py b/mypy/semanal_main.py
@@ -380,7 +380,7 @@ def check_type_arguments(graph: Graph, scc: list[str], errors: Errors) -> None:
         analyzer = TypeArgumentAnalyzer(
             errors,
             state.options,
-            is_typeshed_file(state.options.abs_custom_typeshed_dir, state.path or ""),
+            state.tree.is_typeshed_file(state.options),
             state.manager.semantic_analyzer.named_type,
         )
         with state.wrap_context():

Original file line number	Diff line number	Diff line change
`@@ -380,7 +380,7 @@ def check_type_arguments(graph: Graph, scc: list[str], errors: Errors) -> None:`
`380`	`380`	`analyzer = TypeArgumentAnalyzer(`
`381`	`381`	`errors,`
`382`	`382`	`state.options,`
`383`		`- is_typeshed_file(state.options.abs_custom_typeshed_dir, state.path or ""),`
	`383`	`+ state.tree.is_typeshed_file(state.options),`
`384`	`384`	`state.manager.semantic_analyzer.named_type,`
`385`	`385`	`)`
`386`	`386`	`with state.wrap_context():`