[mypyc] Optimize type(x) and x.__class__ (#19691)

Using `Py_TYPE` avoids a function call. Also specialize `x.__class__`
for instances of native classes -- it's treated as equivalent to
`type(x)`. This is not generally possible, so only do it for native
instances where we can make more assumptions.

Author: JukkaL

mypyc/irbuild/expression.py

@@ -69,6 +69,7 @@
     Value,
 )
 from mypyc.ir.rtypes import (
+    RInstance,
     RTuple,
     bool_rprimitive,
     int_rprimitive,
@@ -226,6 +227,11 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
     ):
         return builder.primitive_op(name_op, [obj], expr.line)
 
+    if isinstance(obj.type, RInstance) and expr.name == "__class__":
+        # A non-native class could override "__class__" using "__getattribute__", so
+        # only apply to RInstance types.
+        return builder.primitive_op(type_op, [obj], expr.line)
+
     # Special case: for named tuples transform attribute access to faster index access.
     typ = get_proper_type(builder.types.get(expr.expr))
     if isinstance(typ, TupleType) and typ.partial_fallback.type.is_named_tuple:

mypyc/lib-rt/CPy.h

@@ -877,6 +877,12 @@ static inline bool CPy_TypeCheck(PyObject *o, PyObject *type) {
     return PyObject_TypeCheck(o, (PyTypeObject *)type);
 }
 
+static inline PyObject *CPy_TYPE(PyObject *obj) {
+    PyObject *result = (PyObject *)Py_TYPE(obj);
+    Py_INCREF(result);
+    return result;
+}
+
 PyObject *CPy_CalculateMetaclass(PyObject *type, PyObject *o);
 PyObject *CPy_GetCoro(PyObject *obj);
 PyObject *CPyIter_Send(PyObject *iter, PyObject *val);

mypyc/primitives/misc_ops.py

@@ -214,7 +214,7 @@
 type_op = function_op(
     name="builtins.type",
     arg_types=[object_rprimitive],
-    c_function_name="PyObject_Type",
+    c_function_name="CPy_TYPE",
     return_type=object_rprimitive,
     error_kind=ERR_NEVER,
 )

mypyc/test-data/fixtures/ir.py

@@ -39,6 +39,7 @@ def __pow__(self, other: T_contra, modulo: _M) -> T_co: ...
 ]
 
 class object:
+    __class__: type
     def __init__(self) -> None: pass
     def __eq__(self, x: object) -> bool: pass
     def __ne__(self, x: object) -> bool: pass

mypyc/test-data/irbuild-classes.test

@@ -1516,3 +1516,45 @@ L0:
     r0 = CPy_GetName(t)
     r1 = cast(str, r0)
     return r1
+
+[case testTypeOfObject]
+class C: pass
+class D(C): pass
+
+def generic_type(x: object) -> type[object]:
+    return type(x)
+
+def generic_class(x: object) -> type[object]:
+    return x.__class__
+
+def native_type(x: C) -> type[object]:
+    return type(x)
+
+def native_class(x: C) -> type[object]:
+    return x.__class__
+[out]
+def generic_type(x):
+    x, r0 :: object
+L0:
+    r0 = CPy_TYPE(x)
+    return r0
+def generic_class(x):
+    x :: object
+    r0 :: str
+    r1 :: object
+L0:
+    r0 = '__class__'
+    r1 = CPyObject_GetAttr(x, r0)
+    return r1
+def native_type(x):
+    x :: __main__.C
+    r0 :: object
+L0:
+    r0 = CPy_TYPE(x)
+    return r0
+def native_class(x):
+    x :: __main__.C
+    r0 :: object
+L0:
+    r0 = CPy_TYPE(x)
+    return r0

mypyc/test-data/irbuild-try.test

@@ -412,7 +412,7 @@ def foo(x):
     r36 :: bit
 L0:
     r0 = PyObject_Vectorcall(x, 0, 0, 0)
-    r1 = PyObject_Type(r0)
+    r1 = CPy_TYPE(r0)
     r2 = '__exit__'
     r3 = CPyObject_GetAttr(r1, r2)
     r4 = '__enter__'

mypyc/test-data/run-classes.test

@@ -3163,3 +3163,49 @@ def foo(): pass
 class Thing:
     def __init__(self):
         self.__name__ = "xyz"
+
+[case testTypeOfObject]
+from typing import Any
+
+from dynamic import Dyn
+
+class Foo: pass
+class Bar(Foo): pass
+
+def generic_type(x) -> type[object]:
+    return x.__class__
+
+def test_built_in_type() -> None:
+    i: Any = int
+    l: Any = list
+    assert type(i()) is i().__class__
+    assert type(i()) is int
+    assert type(l()) is list
+    n = 5
+    assert n.__class__ is i
+
+def test_native_class() -> None:
+    f_any: Any = Foo()
+    b_any: Any = Bar()
+    f: Foo = f_any
+    b: Foo = b_any
+    if int("1"):  # use int("1") to avoid constant folding
+        assert type(f) is Foo
+        assert type(b) is Bar
+    if int("2"):
+        assert f.__class__ is Foo
+        assert b.__class__ is Bar
+    if int("3"):
+        assert f_any.__class__ is Foo
+        assert b_any.__class__ is Bar
+    if int("4"):
+        assert type(f_any) is Foo
+        assert type(b_any) is Bar
+
+def test_python_class() -> None:
+    d = Dyn()
+    assert type(d) is Dyn
+    assert d.__class__ is Dyn
+
+[file dynamic.py]
+class Dyn: pass