[mypyc] Optimize int()/float()/complex() on native classes

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
 

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
 

mypyc/test-data/irbuild-dunders.test

@@ -154,13 +154,21 @@ class C:
     def __abs__(self) -> int:
         return 6
 
+    def __bool__(self) -> bool:
+        return False
+
+    def __complex__(self) -> complex:
+        return 7+1j
+
 def f(c: C) -> None:
     -c
     ~c
     int(c)
     float(c)
     +c
     abs(c)
+    bool(c)
+    complex(c)
 [out]
 def C.__neg__(self):
     self :: __main__.C
@@ -188,19 +196,33 @@ 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, r1, r2 :: object
+L0:
+    r0 = 1j
+    r1 = object 7
+    r2 = PyNumber_Add(r1, r0)
+    return r2
 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