diff --git a/mypy/nodes.py b/mypy/nodes.py index 85bb9ce4a8dee..4a4de9d4503db 100644 --- a/mypy/nodes.py +++ b/mypy/nodes.py @@ -1777,7 +1777,7 @@ class NameExpr(RefExpr): def __init__(self, name: str) -> None: super().__init__() - self.name = name # Name referred to (may be qualified) + self.name = name # Name referred to # Is this a l.h.s. of a special form assignment like typed dict or type variable? self.is_special_form = False diff --git a/mypyc/irbuild/specialize.py b/mypyc/irbuild/specialize.py index 06babd2f7e1a9..e62350778f548 100644 --- a/mypyc/irbuild/specialize.py +++ b/mypyc/irbuild/specialize.py @@ -157,14 +157,20 @@ def translate_globals(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Va @specialize_function("builtins.abs") -def translate_abs(builder: IRBuilder, expr: CallExpr, callee: RefExpr) -> Value | None: - """Specialize calls on native classes that implement __abs__.""" - if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS]: +@specialize_function("builtins.int") +@specialize_function("builtins.float") +@specialize_function("builtins.complex") +def translate_builtins_with_unary_dunder( + builder: IRBuilder, expr: CallExpr, callee: RefExpr +) -> Value | None: + """Specialize calls on native classes that implement the associated dunder.""" + if len(expr.args) == 1 and expr.arg_kinds == [ARG_POS] and isinstance(callee, NameExpr): arg = expr.args[0] arg_typ = builder.node_type(arg) - if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method("__abs__"): + method = f"__{callee.name}__" + if isinstance(arg_typ, RInstance) and arg_typ.class_ir.has_method(method): obj = builder.accept(arg) - return builder.gen_method_call(obj, "__abs__", [], None, expr.line) + return builder.gen_method_call(obj, method, [], None, expr.line) return None diff --git a/mypyc/test-data/irbuild-dunders.test b/mypyc/test-data/irbuild-dunders.test index 24e7089133542..82f04dcdf6870 100644 --- a/mypyc/test-data/irbuild-dunders.test +++ b/mypyc/test-data/irbuild-dunders.test @@ -154,6 +154,12 @@ class C: def __abs__(self) -> int: return 6 + def __bool__(self) -> bool: + return False + + def __complex__(self) -> complex: + return 7j + def f(c: C) -> None: -c ~c @@ -161,6 +167,8 @@ def f(c: C) -> None: float(c) +c abs(c) + bool(c) + complex(c) [out] def C.__neg__(self): self :: __main__.C @@ -188,19 +196,31 @@ def C.__abs__(self): self :: __main__.C L0: return 12 +def C.__bool__(self): + self :: __main__.C +L0: + return 0 +def C.__complex__(self): + self :: __main__.C + r0 :: object +L0: + r0 = 7j + return r0 def f(c): c :: __main__.C - r0, r1 :: int - r2, r3, r4, r5 :: object - r6, r7 :: int + r0, r1, r2 :: int + r3 :: float + r4, r5 :: int + r6 :: bool + r7 :: object L0: r0 = c.__neg__() r1 = c.__invert__() - r2 = load_address PyLong_Type - r3 = PyObject_CallFunctionObjArgs(r2, c, 0) - r4 = load_address PyFloat_Type - r5 = PyObject_CallFunctionObjArgs(r4, c, 0) - r6 = c.__pos__() - r7 = c.__abs__() + r2 = c.__int__() + r3 = c.__float__() + r4 = c.__pos__() + r5 = c.__abs__() + r6 = c.__bool__() + r7 = c.__complex__() return 1