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%.

Author: JukkaL

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 @@ def typeddict_callable_from_context(
         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 @@ def check_typeddict_call_with_kwargs(
                 # 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,

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 @@ def serialize(self) -> JsonDict:
             "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 @@ def deserialize(cls, data: JsonDict) -> FuncDef:
         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 @@ def accept(self, visitor: ExpressionVisitor[T]) -> T:
         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 @@ def is_star(self) -> bool:
         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):

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 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
     # 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 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
             # 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 @@ def handle_partial_with_callee(ctx: mypy.plugin.FunctionContext, callee: Type) -
     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

mypy/semanal.py

@@ -1041,7 +1041,7 @@ def analyze_func_def(self, defn: FuncDef) -> None:
         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):

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 @@ def serialize(self) -> JsonDict:
         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 @@ def deserialize(cls, data: JsonDict) -> Parameters:
         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 @@ def with_unpacked_kwargs(self) -> NormalizedCallableType:
         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 @@ def serialize(self) -> JsonDict:
         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 @@ def deserialize(cls, data: JsonDict) -> CallableType:
         # 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"]),

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 generate_traceback_code(
 
 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]:

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 @@ def serialize(self) -> JsonDict:
         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 @@ def deserialize(cls, data: JsonDict, ctx: DeserMaps) -> RuntimeArg:
         return RuntimeArg(
             data["name"],
             deserialize_type(data["type"], ctx),
-            ArgKind(data["kind"]),
+            ArgKind.by_value(data["kind"]),
             data["pos_only"],
         )
 
@@ -394,12 +404,12 @@ def all_values_full(args: list[Register], blocks: list[BasicBlock]) -> list[Valu
 
 
 _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 @@ def get_text_signature(fn: FuncIR, *, bound: bool = False) -> str | None:
     # 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")):

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 @@ def generate_dispatch_glue_native_function(
     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 @@ def add_register_method_to_callable_class(builder: IRBuilder, fn_info: FuncInfo)
     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))