Micro-optimization: Make ArgKind a regular class instead of enum Mypyc doesn't generate very efficient code for enums, so switch to a regular class. We can later revert the change if/when we can improve enum support in mypyc. Operations related to ArgKind were pretty prominent in the op trace log (#19457). By itself this improves performance by ~1.7%, based on `perf_compare.py`, which is significant: ``` master 4.168s (0.0%) | stdev 0.037s HEAD 4.098s (-1.7%) | stdev 0.028s ``` This is a part of a set of micro-optimizations that improve self check performance by ~5.5%.
diff --git a/mypy/checkexpr.py b/mypy/checkexpr.py index 8223ccf..a8740d4 100644 --- a/mypy/checkexpr.py +++ b/mypy/checkexpr.py
@@ -34,6 +34,7 @@ from mypy.messages import MessageBuilder, format_type from mypy.nodes import ( ARG_NAMED, + ARG_NAMED_OPT, ARG_POS, ARG_STAR, ARG_STAR2, @@ -1000,7 +1001,7 @@ return CallableType( list(callee.items.values()), [ - ArgKind.ARG_NAMED if name in callee.required_keys else ArgKind.ARG_NAMED_OPT + ARG_NAMED if name in callee.required_keys else ARG_NAMED_OPT for name in callee.items ], list(callee.items.keys()), @@ -1074,7 +1075,7 @@ # TypedDict. This is a bit arbitrary, but in most cases will work better than # trying to infer a union or a join. [args[0] for args in kwargs.values()], - [ArgKind.ARG_NAMED] * len(kwargs), + [ARG_NAMED] * len(kwargs), context, list(kwargs.keys()), None,
diff --git a/mypy/nodes.py b/mypy/nodes.py index fc2656c..d61e112 100644 --- a/mypy/nodes.py +++ b/mypy/nodes.py
@@ -6,8 +6,18 @@ from abc import abstractmethod from collections import defaultdict from collections.abc import Iterator, Sequence -from enum import Enum, unique -from typing import TYPE_CHECKING, Any, Callable, Final, Optional, TypeVar, Union, cast +from typing import ( + TYPE_CHECKING, + Any, + Callable, + ClassVar, + Final, + Optional, + TypeVar, + Union, + cast, + final, +) from typing_extensions import TypeAlias as _TypeAlias, TypeGuard from mypy_extensions import trait @@ -873,7 +883,7 @@ "name": self._name, "fullname": self._fullname, "arg_names": self.arg_names, - "arg_kinds": [int(x.value) for x in self.arg_kinds], + "arg_kinds": [x.value for x in self.arg_kinds], "type": None if self.type is None else self.type.serialize(), "flags": get_flags(self, FUNCDEF_FLAGS), "abstract_status": self.abstract_status, @@ -904,7 +914,7 @@ set_flags(ret, data["flags"]) # NOTE: ret.info is set in the fixup phase. ret.arg_names = data["arg_names"] - ret.arg_kinds = [ArgKind(x) for x in data["arg_kinds"]] + ret.arg_kinds = [ArgKind.by_value(x) for x in data["arg_kinds"]] ret.abstract_status = data["abstract_status"] ret.dataclass_transform_spec = ( DataclassTransformSpec.deserialize(data["dataclass_transform_spec"]) @@ -1963,21 +1973,35 @@ return visitor.visit_member_expr(self) -# Kinds of arguments -@unique -class ArgKind(Enum): - # Positional argument - ARG_POS = 0 - # Positional, optional argument (functions only, not calls) - ARG_OPT = 1 - # *arg argument - ARG_STAR = 2 - # Keyword argument x=y in call, or keyword-only function arg - ARG_NAMED = 3 - # **arg argument - ARG_STAR2 = 4 - # In an argument list, keyword-only and also optional - ARG_NAMED_OPT = 5 +@final +class ArgKind: + """Kinds of arguments. + + NOTE: This isn't an enum due to mypyc performance limitations. + """ + + _sealed: ClassVar[bool] = False # Hack to ensure enum-like behavior + + def __init__(self, name: str, value: int) -> None: + assert not ArgKind._sealed + self.name: Final = name + self.value: Final = value + + @staticmethod + def by_value(value: int) -> ArgKind: + if value == ARG_POS.value: + return ARG_POS + elif value == ARG_OPT.value: + return ARG_OPT + elif value == ARG_STAR.value: + return ARG_STAR + elif value == ARG_NAMED.value: + return ARG_NAMED + elif value == ARG_STAR2.value: + return ARG_STAR2 + else: + assert value == ARG_NAMED_OPT.value + return ARG_NAMED_OPT def is_positional(self, star: bool = False) -> bool: return self == ARG_POS or self == ARG_OPT or (star and self == ARG_STAR) @@ -1995,12 +2019,29 @@ return self == ARG_STAR or self == ARG_STAR2 -ARG_POS: Final = ArgKind.ARG_POS -ARG_OPT: Final = ArgKind.ARG_OPT -ARG_STAR: Final = ArgKind.ARG_STAR -ARG_NAMED: Final = ArgKind.ARG_NAMED -ARG_STAR2: Final = ArgKind.ARG_STAR2 -ARG_NAMED_OPT: Final = ArgKind.ARG_NAMED_OPT +# Positional argument +ARG_POS: Final = ArgKind("ARG_POS", 0) +# Positional, optional argument (functions only, not calls) +ARG_OPT: Final = ArgKind("ARG_OPT", 1) +# *arg argument +ARG_STAR: Final = ArgKind("ARG_STAR", 2) +# Keyword argument x=y in call, or keyword-only function arg +ARG_NAMED: Final = ArgKind("ARG_NAMED", 3) +# **arg argument +ARG_STAR2: Final = ArgKind("ARG_STAR2", 4) +# In an argument list, keyword-only and also optional +ARG_NAMED_OPT: Final = ArgKind("ARG_NAMED_OPT", 5) + +ArgKind._sealed = True # Make sure no new ArgKinds can be created + +ALL_ARG_KINDS: Final[tuple[ArgKind, ...]] = ( + ARG_POS, + ARG_OPT, + ARG_STAR, + ARG_NAMED, + ARG_STAR2, + ARG_NAMED_OPT, +) class CallExpr(Expression):
diff --git a/mypy/plugins/functools.py b/mypy/plugins/functools.py index c8b370f..ef62cf4 100644 --- a/mypy/plugins/functools.py +++ b/mypy/plugins/functools.py
@@ -10,10 +10,13 @@ from mypy.argmap import map_actuals_to_formals from mypy.erasetype import erase_typevars from mypy.nodes import ( + ARG_NAMED, + ARG_NAMED_OPT, + ARG_OPT, ARG_POS, + ARG_STAR, ARG_STAR2, SYMBOL_FUNCBASE_TYPES, - ArgKind, Argument, CallExpr, NameExpr, @@ -217,11 +220,7 @@ # special_sig="partial" allows omission of args/kwargs typed with ParamSpec defaulted = fn_type.copy_modified( arg_kinds=[ - ( - ArgKind.ARG_OPT - if k == ArgKind.ARG_POS - else (ArgKind.ARG_NAMED_OPT if k == ArgKind.ARG_NAMED else k) - ) + (ARG_OPT if k == ARG_POS else (ARG_NAMED_OPT if k == ARG_NAMED else k)) for k in fn_type.arg_kinds ], ret_type=ret_type, @@ -284,19 +283,19 @@ # true when PEP 646 things are happening. See testFunctoolsPartialTypeVarTuple arg_type = fn_type.arg_types[i] - if not actuals or fn_type.arg_kinds[i] in (ArgKind.ARG_STAR, ArgKind.ARG_STAR2): + if not actuals or fn_type.arg_kinds[i] in (ARG_STAR, ARG_STAR2): partial_kinds.append(fn_type.arg_kinds[i]) partial_types.append(arg_type) partial_names.append(fn_type.arg_names[i]) else: assert actuals - if any(actual_arg_kinds[j] in (ArgKind.ARG_POS, ArgKind.ARG_STAR) for j in actuals): + if any(actual_arg_kinds[j] in (ARG_POS, ARG_STAR) for j in actuals): # Don't add params for arguments passed positionally continue # Add defaulted params for arguments passed via keyword kind = actual_arg_kinds[actuals[0]] - if kind == ArgKind.ARG_NAMED or kind == ArgKind.ARG_STAR2: - kind = ArgKind.ARG_NAMED_OPT + if kind == ARG_NAMED or kind == ARG_STAR2: + kind = ARG_NAMED_OPT partial_kinds.append(kind) partial_types.append(arg_type) partial_names.append(fn_type.arg_names[i]) @@ -322,9 +321,9 @@ if partially_applied.param_spec(): assert ret.extra_attrs is not None # copy_with_extra_attr above ensures this attrs = ret.extra_attrs.copy() - if ArgKind.ARG_STAR in actual_arg_kinds: + if ARG_STAR in actual_arg_kinds: attrs.immutable.add("__mypy_partial_paramspec_args_bound") - if ArgKind.ARG_STAR2 in actual_arg_kinds: + if ARG_STAR2 in actual_arg_kinds: attrs.immutable.add("__mypy_partial_paramspec_kwargs_bound") ret.extra_attrs = attrs return ret
diff --git a/mypy/semanal.py b/mypy/semanal.py index 01b7f49..14604f9 100644 --- a/mypy/semanal.py +++ b/mypy/semanal.py
@@ -1041,7 +1041,7 @@ self.pop_type_args(defn.type_args) def remove_unpack_kwargs(self, defn: FuncDef, typ: CallableType) -> CallableType: - if not typ.arg_kinds or typ.arg_kinds[-1] is not ArgKind.ARG_STAR2: + if not typ.arg_kinds or typ.arg_kinds[-1] is not ARG_STAR2: return typ last_type = typ.arg_types[-1] if not isinstance(last_type, UnpackType):
diff --git a/mypy/types.py b/mypy/types.py index 05b02ac..51f9c1b 100644 --- a/mypy/types.py +++ b/mypy/types.py
@@ -22,6 +22,8 @@ import mypy.nodes from mypy.bogus_type import Bogus from mypy.nodes import ( + ARG_NAMED, + ARG_NAMED_OPT, ARG_POS, ARG_STAR, ARG_STAR2, @@ -1782,7 +1784,7 @@ return { ".class": "Parameters", "arg_types": [t.serialize() for t in self.arg_types], - "arg_kinds": [int(x.value) for x in self.arg_kinds], + "arg_kinds": [x.value for x in self.arg_kinds], "arg_names": self.arg_names, "variables": [tv.serialize() for tv in self.variables], "imprecise_arg_kinds": self.imprecise_arg_kinds, @@ -1793,7 +1795,7 @@ assert data[".class"] == "Parameters" return Parameters( [deserialize_type(t) for t in data["arg_types"]], - [ArgKind(x) for x in data["arg_kinds"]], + [ArgKind.by_value(x) for x in data["arg_kinds"]], data["arg_names"], variables=[cast(TypeVarLikeType, deserialize_type(v)) for v in data["variables"]], imprecise_arg_kinds=data["imprecise_arg_kinds"], @@ -2162,7 +2164,7 @@ last_type = get_proper_type(self.arg_types[-1]) assert isinstance(last_type, TypedDictType) extra_kinds = [ - ArgKind.ARG_NAMED if name in last_type.required_keys else ArgKind.ARG_NAMED_OPT + ARG_NAMED if name in last_type.required_keys else ARG_NAMED_OPT for name in last_type.items ] new_arg_kinds = self.arg_kinds[:-1] + extra_kinds @@ -2283,7 +2285,7 @@ return { ".class": "CallableType", "arg_types": [t.serialize() for t in self.arg_types], - "arg_kinds": [int(x.value) for x in self.arg_kinds], + "arg_kinds": [x.value for x in self.arg_kinds], "arg_names": self.arg_names, "ret_type": self.ret_type.serialize(), "fallback": self.fallback.serialize(), @@ -2307,7 +2309,7 @@ # TODO: Set definition to the containing SymbolNode? return CallableType( [deserialize_type(t) for t in data["arg_types"]], - [ArgKind(x) for x in data["arg_kinds"]], + [ArgKind.by_value(x) for x in data["arg_kinds"]], data["arg_names"], deserialize_type(data["ret_type"]), Instance.deserialize(data["fallback"]),
diff --git a/mypyc/codegen/emitwrapper.py b/mypyc/codegen/emitwrapper.py index cd16842..ec91ff4 100644 --- a/mypyc/codegen/emitwrapper.py +++ b/mypyc/codegen/emitwrapper.py
@@ -14,7 +14,16 @@ from collections.abc import Sequence -from mypy.nodes import ARG_NAMED, ARG_NAMED_OPT, ARG_OPT, ARG_POS, ARG_STAR, ARG_STAR2, ArgKind +from mypy.nodes import ( + ALL_ARG_KINDS, + ARG_NAMED, + ARG_NAMED_OPT, + ARG_OPT, + ARG_POS, + ARG_STAR, + ARG_STAR2, + ArgKind, +) from mypy.operators import op_methods_to_symbols, reverse_op_method_names, reverse_op_methods from mypyc.codegen.emit import AssignHandler, Emitter, ErrorHandler, GotoHandler, ReturnHandler from mypyc.common import ( @@ -88,7 +97,7 @@ def make_arg_groups(args: list[RuntimeArg]) -> dict[ArgKind, list[RuntimeArg]]: """Group arguments by kind.""" - return {k: [arg for arg in args if arg.kind == k] for k in ArgKind} + return {k: [arg for arg in args if arg.kind == k] for k in ALL_ARG_KINDS} def reorder_arg_groups(groups: dict[ArgKind, list[RuntimeArg]]) -> list[RuntimeArg]:
diff --git a/mypyc/ir/func_ir.py b/mypyc/ir/func_ir.py index 881ac59..83c0712 100644 --- a/mypyc/ir/func_ir.py +++ b/mypyc/ir/func_ir.py
@@ -6,7 +6,17 @@ from collections.abc import Sequence from typing import Final -from mypy.nodes import ARG_POS, ArgKind, Block, FuncDef +from mypy.nodes import ( + ARG_NAMED, + ARG_NAMED_OPT, + ARG_OPT, + ARG_POS, + ARG_STAR, + ARG_STAR2, + ArgKind, + Block, + FuncDef, +) from mypyc.common import BITMAP_BITS, JsonDict, bitmap_name, get_id_from_name, short_id_from_name from mypyc.ir.ops import ( Assign, @@ -60,7 +70,7 @@ return { "name": self.name, "type": self.type.serialize(), - "kind": int(self.kind.value), + "kind": self.kind.value, "pos_only": self.pos_only, } @@ -69,7 +79,7 @@ return RuntimeArg( data["name"], deserialize_type(data["type"], ctx), - ArgKind(data["kind"]), + ArgKind.by_value(data["kind"]), data["pos_only"], ) @@ -394,12 +404,12 @@ _ARG_KIND_TO_INSPECT: Final = { - ArgKind.ARG_POS: inspect.Parameter.POSITIONAL_OR_KEYWORD, - ArgKind.ARG_OPT: inspect.Parameter.POSITIONAL_OR_KEYWORD, - ArgKind.ARG_STAR: inspect.Parameter.VAR_POSITIONAL, - ArgKind.ARG_NAMED: inspect.Parameter.KEYWORD_ONLY, - ArgKind.ARG_STAR2: inspect.Parameter.VAR_KEYWORD, - ArgKind.ARG_NAMED_OPT: inspect.Parameter.KEYWORD_ONLY, + ARG_POS: inspect.Parameter.POSITIONAL_OR_KEYWORD, + ARG_OPT: inspect.Parameter.POSITIONAL_OR_KEYWORD, + ARG_STAR: inspect.Parameter.VAR_POSITIONAL, + ARG_NAMED: inspect.Parameter.KEYWORD_ONLY, + ARG_STAR2: inspect.Parameter.VAR_KEYWORD, + ARG_NAMED_OPT: inspect.Parameter.KEYWORD_ONLY, } # Sentinel indicating a value that cannot be represented in a text signature. @@ -418,7 +428,7 @@ # currently sees 'self' as being positional-or-keyword and '__x' as positional-only. pos_only_idx = -1 for idx, arg in enumerate(sig.args): - if arg.pos_only and arg.kind in (ArgKind.ARG_POS, ArgKind.ARG_OPT): + if arg.pos_only and arg.kind in (ARG_POS, ARG_OPT): pos_only_idx = idx for idx, arg in enumerate(sig.args): if arg.name.startswith(("__bitmap", "__mypyc")):
diff --git a/mypyc/irbuild/function.py b/mypyc/irbuild/function.py index 90506ad..ac4b844 100644 --- a/mypyc/irbuild/function.py +++ b/mypyc/irbuild/function.py
@@ -17,6 +17,8 @@ from typing import NamedTuple from mypy.nodes import ( + ARG_OPT, + ARG_POS, ArgKind, ClassDef, Decorator, @@ -907,7 +909,7 @@ decl = builder.mapper.func_to_decl[fitem] arg_info = get_args(builder, decl.sig.args, line) args = [callable_class] + arg_info.args - arg_kinds = [ArgKind.ARG_POS] + arg_info.arg_kinds + arg_kinds = [ARG_POS] + arg_info.arg_kinds arg_names = arg_info.arg_names arg_names.insert(0, "self") ret_val = builder.builder.call(callable_class_decl, args, arg_kinds, arg_names, line) @@ -935,7 +937,7 @@ line = -1 with builder.enter_method(fn_info.callable_class.ir, "register", object_rprimitive): cls_arg = builder.add_argument("cls", object_rprimitive) - func_arg = builder.add_argument("func", object_rprimitive, ArgKind.ARG_OPT) + func_arg = builder.add_argument("func", object_rprimitive, ARG_OPT) ret_val = builder.call_c(register_function, [builder.self(), cls_arg, func_arg], line) builder.add(Return(ret_val, line))