Speed up the default plugin (#19385)

Fix two kinds of inefficiency in the default plugin:
* Pre-calculate various set objects, since set construction is pretty
slow
* Nested imports are pretty slow, so avoid doing them in the fast path

I also had to refactor things a little in order to optimize the nested
imports.

This speeds up self check by about 2.2%. The default plugin is called a
lot.

Author: JukkaL

mypy/plugins/constants.py

@@ -0,0 +1,20 @@
+"""Constant definitions for plugins kept here to help with import cycles."""
+
+from typing import Final
+
+from mypy.semanal_enum import ENUM_BASES
+
+SINGLEDISPATCH_TYPE: Final = "functools._SingleDispatchCallable"
+SINGLEDISPATCH_REGISTER_METHOD: Final = f"{SINGLEDISPATCH_TYPE}.register"
+SINGLEDISPATCH_CALLABLE_CALL_METHOD: Final = f"{SINGLEDISPATCH_TYPE}.__call__"
+SINGLEDISPATCH_REGISTER_RETURN_CLASS: Final = "_SingleDispatchRegisterCallable"
+SINGLEDISPATCH_REGISTER_CALLABLE_CALL_METHOD: Final = (
+    f"functools.{SINGLEDISPATCH_REGISTER_RETURN_CLASS}.__call__"
+)
+
+ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
+    f"{prefix}._name_" for prefix in ENUM_BASES
+}
+ENUM_VALUE_ACCESS: Final = {f"{prefix}.value" for prefix in ENUM_BASES} | {
+    f"{prefix}._value_" for prefix in ENUM_BASES
+}

mypy/plugins/default.py

@@ -15,6 +15,7 @@
     MethodSigContext,
     Plugin,
 )
+from mypy.plugins import constants
 from mypy.plugins.common import try_getting_str_literals
 from mypy.subtypes import is_subtype
 from mypy.typeops import is_literal_type_like, make_simplified_union
@@ -36,70 +37,79 @@
     get_proper_types,
 )
 
+TD_SETDEFAULT_NAMES: Final = {n + ".setdefault" for n in TPDICT_FB_NAMES}
+TD_POP_NAMES: Final = {n + ".pop" for n in TPDICT_FB_NAMES}
+
+TD_UPDATE_METHOD_NAMES: Final = (
+    {n + ".update" for n in TPDICT_FB_NAMES}
+    | {n + ".__or__" for n in TPDICT_FB_NAMES}
+    | {n + ".__ror__" for n in TPDICT_FB_NAMES}
+    | {n + ".__ior__" for n in TPDICT_FB_NAMES}
+)
+
 
 class DefaultPlugin(Plugin):
     """Type checker plugin that is enabled by default."""
 
     def get_function_hook(self, fullname: str) -> Callable[[FunctionContext], Type] | None:
-        from mypy.plugins import ctypes, enums, singledispatch
-
         if fullname == "_ctypes.Array":
+            from mypy.plugins import ctypes
+
             return ctypes.array_constructor_callback
         elif fullname == "functools.singledispatch":
+            from mypy.plugins import singledispatch
+
             return singledispatch.create_singledispatch_function_callback
         elif fullname == "functools.partial":
             import mypy.plugins.functools
 
             return mypy.plugins.functools.partial_new_callback
         elif fullname == "enum.member":
+            from mypy.plugins import enums
+
             return enums.enum_member_callback
 
         return None
 
     def get_function_signature_hook(
         self, fullname: str
     ) -> Callable[[FunctionSigContext], FunctionLike] | None:
-        from mypy.plugins import attrs, dataclasses
-
         if fullname in ("attr.evolve", "attrs.evolve", "attr.assoc", "attrs.assoc"):
+            from mypy.plugins import attrs
+
             return attrs.evolve_function_sig_callback
         elif fullname in ("attr.fields", "attrs.fields"):
+            from mypy.plugins import attrs
+
             return attrs.fields_function_sig_callback
         elif fullname == "dataclasses.replace":
+            from mypy.plugins import dataclasses
+
             return dataclasses.replace_function_sig_callback
         return None
 
     def get_method_signature_hook(
         self, fullname: str
     ) -> Callable[[MethodSigContext], FunctionLike] | None:
-        from mypy.plugins import ctypes, singledispatch
-
         if fullname == "typing.Mapping.get":
             return typed_dict_get_signature_callback
-        elif fullname in {n + ".setdefault" for n in TPDICT_FB_NAMES}:
+        elif fullname in TD_SETDEFAULT_NAMES:
             return typed_dict_setdefault_signature_callback
-        elif fullname in {n + ".pop" for n in TPDICT_FB_NAMES}:
+        elif fullname in TD_POP_NAMES:
             return typed_dict_pop_signature_callback
         elif fullname == "_ctypes.Array.__setitem__":
-            return ctypes.array_setitem_callback
-        elif fullname == singledispatch.SINGLEDISPATCH_CALLABLE_CALL_METHOD:
-            return singledispatch.call_singledispatch_function_callback
+            from mypy.plugins import ctypes
 
-        typed_dict_updates = set()
-        for n in TPDICT_FB_NAMES:
-            typed_dict_updates.add(n + ".update")
-            typed_dict_updates.add(n + ".__or__")
-            typed_dict_updates.add(n + ".__ror__")
-            typed_dict_updates.add(n + ".__ior__")
+            return ctypes.array_setitem_callback
+        elif fullname == constants.SINGLEDISPATCH_CALLABLE_CALL_METHOD:
+            from mypy.plugins import singledispatch
 
-        if fullname in typed_dict_updates:
+            return singledispatch.call_singledispatch_function_callback
+        elif fullname in TD_UPDATE_METHOD_NAMES:
             return typed_dict_update_signature_callback
-
         return None
 
     def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | None:
-        from mypy.plugins import ctypes, singledispatch
-
         if fullname == "typing.Mapping.get":
             return typed_dict_get_callback
         elif fullname == "builtins.int.__pow__":
@@ -117,12 +127,20 @@ def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | No
         elif fullname in {n + ".__delitem__" for n in TPDICT_FB_NAMES}:
             return typed_dict_delitem_callback
         elif fullname == "_ctypes.Array.__getitem__":
+            from mypy.plugins import ctypes
+
             return ctypes.array_getitem_callback
         elif fullname == "_ctypes.Array.__iter__":
+            from mypy.plugins import ctypes
+
             return ctypes.array_iter_callback
-        elif fullname == singledispatch.SINGLEDISPATCH_REGISTER_METHOD:
+        elif fullname == constants.SINGLEDISPATCH_REGISTER_METHOD:
+            from mypy.plugins import singledispatch
+
             return singledispatch.singledispatch_register_callback
-        elif fullname == singledispatch.REGISTER_CALLABLE_CALL_METHOD:
+        elif fullname == constants.SINGLEDISPATCH_REGISTER_CALLABLE_CALL_METHOD:
+            from mypy.plugins import singledispatch
+
             return singledispatch.call_singledispatch_function_after_register_argument
         elif fullname == "functools.partial.__call__":
             import mypy.plugins.functools
@@ -131,15 +149,21 @@ def get_method_hook(self, fullname: str) -> Callable[[MethodContext], Type] | No
         return None
 
     def get_attribute_hook(self, fullname: str) -> Callable[[AttributeContext], Type] | None:
-        from mypy.plugins import ctypes, enums
-
         if fullname == "_ctypes.Array.value":
+            from mypy.plugins import ctypes
+
             return ctypes.array_value_callback
         elif fullname == "_ctypes.Array.raw":
+            from mypy.plugins import ctypes
+
             return ctypes.array_raw_callback
-        elif fullname in enums.ENUM_NAME_ACCESS:
+        elif fullname in constants.ENUM_NAME_ACCESS:
+            from mypy.plugins import enums
+
             return enums.enum_name_callback
-        elif fullname in enums.ENUM_VALUE_ACCESS:
+        elif fullname in constants.ENUM_VALUE_ACCESS:
+            from mypy.plugins import enums
+
             return enums.enum_value_callback
         return None
 

mypy/plugins/enums.py

@@ -14,11 +14,10 @@
 from __future__ import annotations
 
 from collections.abc import Iterable, Sequence
-from typing import Final, TypeVar, cast
+from typing import TypeVar, cast
 
 import mypy.plugin  # To avoid circular imports.
 from mypy.nodes import TypeInfo
-from mypy.semanal_enum import ENUM_BASES
 from mypy.subtypes import is_equivalent
 from mypy.typeops import fixup_partial_type, make_simplified_union
 from mypy.types import (
@@ -31,13 +30,6 @@
     is_named_instance,
 )
 
-ENUM_NAME_ACCESS: Final = {f"{prefix}.name" for prefix in ENUM_BASES} | {
-    f"{prefix}._name_" for prefix in ENUM_BASES
-}
-ENUM_VALUE_ACCESS: Final = {f"{prefix}.value" for prefix in ENUM_BASES} | {
-    f"{prefix}._value_" for prefix in ENUM_BASES
-}
-
 
 def enum_name_callback(ctx: mypy.plugin.AttributeContext) -> Type:
     """This plugin refines the 'name' attribute in enums to act as if

mypy/plugins/singledispatch.py

@@ -1,14 +1,15 @@
 from __future__ import annotations
 
 from collections.abc import Sequence
-from typing import Final, NamedTuple, TypeVar, Union
+from typing import NamedTuple, TypeVar, Union
 from typing_extensions import TypeAlias as _TypeAlias
 
 from mypy.messages import format_type
 from mypy.nodes import ARG_POS, Argument, Block, ClassDef, Context, SymbolTable, TypeInfo, Var
 from mypy.options import Options
 from mypy.plugin import CheckerPluginInterface, FunctionContext, MethodContext, MethodSigContext
 from mypy.plugins.common import add_method_to_class
+from mypy.plugins.constants import SINGLEDISPATCH_REGISTER_RETURN_CLASS
 from mypy.subtypes import is_subtype
 from mypy.types import (
     AnyType,
@@ -33,13 +34,6 @@ class RegisterCallableInfo(NamedTuple):
     singledispatch_obj: Instance
 
 
-SINGLEDISPATCH_TYPE: Final = "functools._SingleDispatchCallable"
-
-SINGLEDISPATCH_REGISTER_METHOD: Final = f"{SINGLEDISPATCH_TYPE}.register"
-
-SINGLEDISPATCH_CALLABLE_CALL_METHOD: Final = f"{SINGLEDISPATCH_TYPE}.__call__"
-
-
 def get_singledispatch_info(typ: Instance) -> SingledispatchTypeVars | None:
     if len(typ.args) == 2:
         return SingledispatchTypeVars(*typ.args)  # type: ignore[arg-type]
@@ -56,16 +50,11 @@ def get_first_arg(args: list[list[T]]) -> T | None:
     return None
 
 
-REGISTER_RETURN_CLASS: Final = "_SingleDispatchRegisterCallable"
-
-REGISTER_CALLABLE_CALL_METHOD: Final = f"functools.{REGISTER_RETURN_CLASS}.__call__"
-
-
 def make_fake_register_class_instance(
     api: CheckerPluginInterface, type_args: Sequence[Type]
 ) -> Instance:
-    defn = ClassDef(REGISTER_RETURN_CLASS, Block([]))
-    defn.fullname = f"functools.{REGISTER_RETURN_CLASS}"
+    defn = ClassDef(SINGLEDISPATCH_REGISTER_RETURN_CLASS, Block([]))
+    defn.fullname = f"functools.{SINGLEDISPATCH_REGISTER_RETURN_CLASS}"
     info = TypeInfo(SymbolTable(), defn, "functools")
     obj_type = api.named_generic_type("builtins.object", []).type
     info.bases = [Instance(obj_type, [])]