Skip to content

Commit

Permalink
[mypyc] Optimize int()/float()/complex() on native classes (#14450)
Browse files Browse the repository at this point in the history
int() and float() calls on native classes can simply call the associated
dunder if the RInstance defines it, no need to load the type and call
it.

bool() calls were already optimized merely a few days ago, but there
wasn't an IRbuild test verifying this so I added one.

---

Follow up to #14422. I saw the PR and
it reminded me that I had this old patch laying around :)
  • Loading branch information
ichard26 authored Jan 18, 2023
1 parent f625602 commit acf26f4
Show file tree
Hide file tree
Showing 3 changed files with 41 additions and 15 deletions.
2 changes: 1 addition & 1 deletion mypy/nodes.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
16 changes: 11 additions & 5 deletions mypyc/irbuild/specialize.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
38 changes: 29 additions & 9 deletions mypyc/test-data/irbuild-dunders.test
Original file line number Diff line number Diff line change
Expand Up @@ -154,13 +154,21 @@ 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
int(c)
float(c)
+c
abs(c)
bool(c)
complex(c)
[out]
def C.__neg__(self):
self :: __main__.C
Expand Down Expand Up @@ -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

0 comments on commit acf26f4

Please sign in to comment.