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))
