[mypyc] Specialize .real and .imag on complex, float, int and bool

mypyc/doc/float_operations.rst

@@ -26,6 +26,15 @@ Operators
 * Comparisons (``==``, ``!=``, ``<``, etc.)
 * Augmented assignment (``x += y``, etc.)
 
+Attributes
+----------
+
+* ``f.real``
+* ``f.imag``
+
+These are also fast for ``complex`` values and for unions such as
+``int | float | complex``.
+
 Functions
 ---------
 

mypyc/doc/int_operations.rst

@@ -143,6 +143,12 @@ and the ``int`` operand is coerced to the native integer type::
       b = 1 - x
       # Similarly, type of "b" is "i64"
 
+Attributes
+----------
+
+* ``n.real`` (``int`` only)
+* ``n.imag`` (``int`` only)
+
 Methods
 -------
 

mypyc/irbuild/expression.py

@@ -54,8 +54,10 @@
     Instance,
     ProperType,
     TupleType,
+    Type,
     TypeOfAny,
     TypeType,
+    UnionType,
     get_proper_type,
 )
 from mypyc.common import (
@@ -73,6 +75,7 @@
     CallC,
     Cast,
     ComparisonOp,
+    Float,
     GetAttr,
     Integer,
     LoadAddress,
@@ -88,14 +91,17 @@
 from mypyc.ir.rtypes import (
     RInstance,
     RTuple,
+    RType,
     RVec,
     bool_rprimitive,
     int64_rprimitive,
     int_rprimitive,
     is_any_int,
+    is_bool_or_bit_rprimitive,
     is_bytearray_rprimitive,
     is_bytes_rprimitive,
     is_fixed_width_rtype,
+    is_float_rprimitive,
     is_int64_rprimitive,
     is_int_rprimitive,
     is_list_rprimitive,
@@ -145,6 +151,12 @@
 )
 from mypyc.primitives.bytes_ops import bytes_slice_op
 from mypyc.primitives.dict_ops import dict_get_item_op, dict_new_op, exact_dict_set_item_op
+from mypyc.primitives.float_ops import (
+    complex_imag_op,
+    complex_real_op,
+    number_imag_op,
+    number_real_op,
+)
 from mypyc.primitives.generic_ops import iter_op, name_op
 from mypyc.primitives.list_ops import list_append_op, list_extend_op, list_slice_op
 from mypyc.primitives.misc_ops import ellipsis_op, get_module_dict_op, new_slice_op, type_op
@@ -303,6 +315,9 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
         # only apply to RInstance types.
         return builder.primitive_op(type_op, [obj], expr.line)
 
+    if expr.name in ("real", "imag") and is_builtin_number_type(builder.types.get(expr.expr)):
+        return transform_real_imag(builder, obj, expr.name == "real", rtype, 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:
@@ -335,6 +350,37 @@ def transform_member_expr(builder: IRBuilder, expr: MemberExpr) -> Value:
     )
 
 
+def is_builtin_number_type(typ: Type | None) -> bool:
+    """Is typ int, bool, float, complex or a union of them (and not a subclass)?"""
+    typ = get_proper_type(typ)
+    if isinstance(typ, UnionType):
+        return all(is_builtin_number_type(item) for item in typ.items)
+    return isinstance(typ, Instance) and typ.type.fullname in (
+        "builtins.int",
+        "builtins.bool",
+        "builtins.float",
+        "builtins.complex",
+    )
+
+
+def transform_real_imag(
+    builder: IRBuilder, obj: Value, is_real: bool, rtype: RType, line: int
+) -> Value:
+    """Get obj.real or obj.imag, where obj is an int, a float or a complex."""
+    if is_tagged(obj.type) or is_bool_or_bit_rprimitive(obj.type):
+        # Unboxed ints are always exact ints, and the real part of a bool is an int
+        return builder.coerce(obj, int_rprimitive, line) if is_real else builder.load_int(0, line)
+    if is_float_rprimitive(obj.type):
+        # Unboxed floats are always exact floats, so these can't be overridden
+        return obj if is_real else Float(0.0, line)
+    if is_float_rprimitive(rtype):
+        op = complex_real_op if is_real else complex_imag_op
+        return builder.primitive_op(op, [obj], line)
+    # The result can be an int, for example with int | float | complex
+    op = number_real_op if is_real else number_imag_op
+    return builder.primitive_op(op, [obj], line, result_type=rtype)
+
+
 def value_borrow_scope(builder: IRBuilder, v: Value) -> int:
     """Compute how long an existing borrowed value can safely be kept alive.
 

mypyc/lib-rt/CPy.h

@@ -640,6 +640,57 @@ CPyTagged CPyFloat_Ceil(double x);
 double CPyFloat_FromTagged(CPyTagged x);
 bool CPyFloat_IsInf(double x);
 bool CPyFloat_IsNaN(double x);
+double CPyComplex_GetPartSlow(PyObject *o, PyObject *name);
+double CPyComplex_LongAsDouble(PyObject *o);
+
+// o.real, where the static type of o is float or complex. Other types,
+// including subclasses that might override the property, use getattr.
+static inline double CPyComplex_Real(PyObject *o) {
+    if (PyComplex_CheckExact(o)) {
+        return ((PyComplexObject *)o)->cval.real;
+    } else if (PyFloat_CheckExact(o)) {
+        return PyFloat_AS_DOUBLE(o);
+    } else if (PyLong_CheckExact(o)) {
+        // An int can be used where a float or complex is expected
+        return CPyComplex_LongAsDouble(o);
+    }
+    return CPyComplex_GetPartSlow(o, mypyc_interned_str.real);
+}
+
+// o.imag, where the static type of o is float or complex
+static inline double CPyComplex_Imag(PyObject *o) {
+    if (PyComplex_CheckExact(o)) {
+        return ((PyComplexObject *)o)->cval.imag;
+    } else if (PyFloat_CheckExact(o) || PyLong_CheckExact(o)) {
+        return 0.0;
+    }
+    return CPyComplex_GetPartSlow(o, mypyc_interned_str.imag);
+}
+
+// o.real, where the static type of o is a union such as int | float | complex,
+// so the result can be an int or a float. Other types use getattr.
+static inline PyObject *CPyNumber_Real(PyObject *o) {
+    if (PyLong_CheckExact(o) || PyFloat_CheckExact(o)) {
+        // int.real and float.real return the object itself
+        Py_INCREF(o);
+        return o;
+    } else if (PyComplex_CheckExact(o)) {
+        return PyFloat_FromDouble(((PyComplexObject *)o)->cval.real);
+    }
+    return PyObject_GetAttr(o, mypyc_interned_str.real);
+}
+
+// o.imag, where the static type of o is a union such as int | float | complex
+static inline PyObject *CPyNumber_Imag(PyObject *o) {
+    if (PyLong_CheckExact(o)) {
+        return PyLong_FromLong(0);
+    } else if (PyFloat_CheckExact(o)) {
+        return PyFloat_FromDouble(0.0);
+    } else if (PyComplex_CheckExact(o)) {
+        return PyFloat_FromDouble(((PyComplexObject *)o)->cval.imag);
+    }
+    return PyObject_GetAttr(o, mypyc_interned_str.imag);
+}
 
 
 // Generic operations (that work with arbitrary types)

mypyc/lib-rt/float_ops.c

@@ -237,3 +237,30 @@ double CPyFloat_Pow(double x, double y) {
     }
     return r;
 }
+
+// Look up o.real or o.imag and convert it to a C double (see CPyComplex_Real)
+double CPyComplex_GetPartSlow(PyObject *o, PyObject *name) {
+    PyObject *part = PyObject_GetAttr(o, name);
+    if (part == NULL) {
+        return CPY_FLOAT_ERROR;
+    }
+    double result = PyFloat_AsDouble(part);
+    if (result == -1.0 && PyErr_Occurred()) {
+        // Same error as unboxing the attribute value would produce
+        CPy_TypeError("float", part);
+        result = CPY_FLOAT_ERROR;
+    }
+    Py_DECREF(part);
+    return result;
+}
+
+// Convert the real part of an int to a C double (see CPyComplex_Real)
+double CPyComplex_LongAsDouble(PyObject *o) {
+    double result = PyLong_AsDouble(o);
+    if (result == -1.0 && PyErr_Occurred()) {
+        // Same error as unboxing the attribute value would produce
+        CPy_TypeError("float", o);
+        return CPY_FLOAT_ERROR;
+    }
+    return result;
+}

mypyc/lib-rt/static_data.c

@@ -54,10 +54,12 @@ intern_strings(void) {
     INTERN_STRING(dispatch_cache, "dispatch_cache");
     INTERN_STRING(endswith, "endswith");
     INTERN_STRING(get_type_hints, "get_type_hints");
+    INTERN_STRING(imag, "imag");
     INTERN_STRING(keys, "keys");
     INTERN_STRING(lower, "lower");
     INTERN_STRING(items, "items");
     INTERN_STRING(join, "join");
+    INTERN_STRING(real, "real");
     INTERN_STRING(register_, "register");
     INTERN_STRING(registry, "registry");
     INTERN_STRING(send, "send");

mypyc/lib-rt/static_data.h

@@ -46,10 +46,12 @@ typedef struct mypyc_interned_str_struct {
     PyObject *dispatch_cache;
     PyObject *endswith;
     PyObject *get_type_hints;
+    PyObject *imag;
     PyObject *keys;
     PyObject *lower;
     PyObject *items;
     PyObject *join;
+    PyObject *real;
     PyObject *register_;
     PyObject *registry;
     PyObject *send;

mypyc/primitives/float_ops.py

@@ -11,7 +11,7 @@
     object_rprimitive,
     str_rprimitive,
 )
-from mypyc.primitives.registry import binary_op, function_op, load_address_op
+from mypyc.primitives.registry import binary_op, custom_primitive_op, function_op, load_address_op
 
 # Get the 'builtins.float' type object.
 load_address_op(name="builtins.float", type=object_rprimitive, src="PyFloat_Type")
@@ -176,3 +176,39 @@
     c_function_name="PyFloat_Check",
     error_kind=ERR_NEVER,
 )
+
+# obj.real, where obj is a float or a complex
+complex_real_op = custom_primitive_op(
+    name="complex.real",
+    arg_types=[object_rprimitive],
+    return_type=float_rprimitive,
+    c_function_name="CPyComplex_Real",
+    error_kind=ERR_MAGIC_OVERLAPPING,
+)
+
+# obj.imag, where obj is a float or a complex
+complex_imag_op = custom_primitive_op(
+    name="complex.imag",
+    arg_types=[object_rprimitive],
+    return_type=float_rprimitive,
+    c_function_name="CPyComplex_Imag",
+    error_kind=ERR_MAGIC_OVERLAPPING,
+)
+
+# obj.real, where obj is an int, a float or a complex (the result can be an int)
+number_real_op = custom_primitive_op(
+    name="number.real",
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPyNumber_Real",
+    error_kind=ERR_MAGIC,
+)
+
+# obj.imag, where obj is an int, a float or a complex (the result can be an int)
+number_imag_op = custom_primitive_op(
+    name="number.imag",
+    arg_types=[object_rprimitive],
+    return_type=object_rprimitive,
+    c_function_name="CPyNumber_Imag",
+    error_kind=ERR_MAGIC,
+)

mypyc/test-data/fixtures/ir.py

@@ -90,6 +90,10 @@ def __le__(self, n: int) -> bool: pass
     def __ge__(self, n: int) -> bool: pass
     def to_bytes(self, length: int, order: str, *, signed: bool = False) -> bytes: pass
     def bit_length(self) -> int: pass
+    @property
+    def real(self) -> int: ...
+    @property
+    def imag(self) -> int: ...
 
 class str:
     @overload
@@ -158,6 +162,10 @@ def __lt__(self, x: float) -> bool: ...
     def __le__(self, x: float) -> bool: ...
     def __gt__(self, x: float) -> bool: ...
     def __ge__(self, x: float) -> bool: ...
+    @property
+    def real(self) -> float: ...
+    @property
+    def imag(self) -> float: ...
 
 class complex:
     def __init__(self, x: object, y: object = None) -> None: pass
@@ -168,6 +176,10 @@ def __rsub__(self, n: float) -> complex: pass
     def __mul__(self, n: complex) -> complex: pass
     def __truediv__(self, n: complex) -> complex: pass
     def __neg__(self) -> complex: pass
+    @property
+    def real(self) -> float: ...
+    @property
+    def imag(self) -> float: ...
 
 class bytes:
     @overload