Implement creating native complex subclasses

Author: msimacek

graalpython/com.oracle.graal.python.cext/src/complexobject.c

@@ -221,7 +221,6 @@ _Py_c_abs(Py_complex z)
     return result;
 }
 
-#if 0 // GraalPy change
 static PyObject *
 complex_subtype_from_c_complex(PyTypeObject *type, Py_complex cval)
 {
@@ -232,7 +231,6 @@ complex_subtype_from_c_complex(PyTypeObject *type, Py_complex cval)
         ((PyComplexObject *)op)->cval = cval;
     return op;
 }
-#endif // GraalPy change
 
 PyObject *
 PyComplex_FromCComplex(Py_complex cval)
@@ -241,8 +239,7 @@ PyComplex_FromCComplex(Py_complex cval)
 	return GraalPyComplex_FromDoubles(cval.real, cval.imag);
 }
 
-#if 0 // GraalPy change
-static PyObject *
+PyObject * // GraalPy change: remove static
 complex_subtype_from_doubles(PyTypeObject *type, double real, double imag)
 {
     Py_complex c;
@@ -251,6 +248,7 @@ complex_subtype_from_doubles(PyTypeObject *type, double real, double imag)
     return complex_subtype_from_c_complex(type, c);
 }
 
+#if 0 // GraalPy change
 PyObject *
 PyComplex_FromDoubles(double real, double imag)
 {
@@ -259,29 +257,33 @@ PyComplex_FromDoubles(double real, double imag)
     c.imag = imag;
     return PyComplex_FromCComplex(c);
 }
+#endif // GraalPy change
 
 double
 PyComplex_RealAsDouble(PyObject *op)
 {
-    if (PyComplex_Check(op)) {
+    // GraalPy change: different implementation
+    if (!points_to_py_handle_space(op) && PyComplex_Check(op)) {
         return ((PyComplexObject *)op)->cval.real;
     }
     else {
-        return PyFloat_AsDouble(op);
+        return GraalPyTruffleComplex_RealAsDouble(op);
     }
 }
 
 double
 PyComplex_ImagAsDouble(PyObject *op)
 {
-    if (PyComplex_Check(op)) {
+    // GraalPy change: different implementation
+    if (!points_to_py_handle_space(op) && PyComplex_Check(op)) {
         return ((PyComplexObject *)op)->cval.imag;
     }
     else {
-        return 0.0;
+        return GraalPyTruffleComplex_ImagAsDouble(op);
     }
 }
 
+#if 0 // GraalPy change
 static PyObject *
 try_complex_special_method(PyObject *op)
 {
@@ -320,6 +322,10 @@ Py_complex
 PyComplex_AsCComplex(PyObject *op)
 {
     // GraalPy change: different implementation
+    /* If op is already of type PyComplex_Type, return its value */
+    if (!points_to_py_handle_space(op) && PyComplex_Check(op)) {
+        return ((PyComplexObject *)op)->cval;
+    }
 	PyObject* parts = GraalPyTruffleComplex_AsCComplex(op);
 	Py_complex result;
 	if(parts != NULL) {

graalpython/com.oracle.graal.python.test/src/tests/cpyext/test_complex.py

@@ -36,9 +36,10 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
+import operator
+
+from . import CPyExtTestCase, CPyExtFunction, unhandled_error_compare, CPyExtType, is_native_object
 
-import sys
-from . import CPyExtTestCase, CPyExtFunction, CPyExtFunctionOutVars, unhandled_error_compare, GRAALPYTHON
 __dir__ = __file__.rpartition("/")[0]
 
 
@@ -84,13 +85,30 @@ class DummyComplexSubclass(complex):
     pass
 
 
+NativeComplexSubclass = CPyExtType(
+    "NativeComplexSubclass",
+    '',
+    struct_base='PyComplexObject base;',
+    tp_base='&PyComplex_Type',
+    tp_new='0',
+    tp_alloc='0',
+    tp_free='0',
+)
+
+
+class ManagedNativeComplexSubclass(NativeComplexSubclass):
+    pass
+
+
 class TestPyComplex(CPyExtTestCase):
 
     test_PyComplex_AsCComplex = CPyExtFunction(
         lambda args: True,
         lambda: (
             (complex(1.0, 2.0), 1.0, 2.0),
             (DummyComplexSubclass(2.0, 3.0), 2.0, 3.0),
+            (NativeComplexSubclass(1.0, 2.0), 1.0, 2.0),
+            (ManagedNativeComplexSubclass(1.0, 2.0), 1.0, 2.0),
         ),
         code='''int isNaN(double d) {
             return d != d;
@@ -122,6 +140,8 @@ class TestPyComplex(CPyExtTestCase):
         lambda: (
             (complex(1.0, 2.0), ),
             (DummyComplexSubclass(2.0, 3.0), ),
+            (NativeComplexSubclass(1.0, 2.0),),
+            (ManagedNativeComplexSubclass(1.0, 2.0),),
         ),
         code='''
         PyObject* wrap_PyComplex_cval(PyObject* obj) {
@@ -143,6 +163,8 @@ class TestPyComplex(CPyExtTestCase):
             (complex(0.0, 2.0), ),
             (complex(1.0, 2.0), ),
             (DummyComplexSubclass(2.0, 3.0), ),
+            (NativeComplexSubclass(1.0, 2.0),),
+            (ManagedNativeComplexSubclass(1.0, 2.0),),
             ("10.0", ),
         ),
         resultspec="f",
@@ -157,6 +179,8 @@ class TestPyComplex(CPyExtTestCase):
             (complex(0.0, 2.0), ),
             (complex(1.0, 2.0), ),
             (DummyComplexSubclass(2.0, 3.0), ),
+            (NativeComplexSubclass(1.0, 2.0),),
+            (ManagedNativeComplexSubclass(1.0, 2.0),),
             ("10.0", ),
         ),
         resultspec="f",
@@ -176,3 +200,35 @@ class TestPyComplex(CPyExtTestCase):
         arguments=["double r", "double i"],
         cmpfunc=unhandled_error_compare
     )
+
+
+class TestNativeComplex:
+    def test_builtins_on_subclass(self):
+        for t in [NativeComplexSubclass, ManagedNativeComplexSubclass]:
+            c = t(2, 3)
+            assert is_native_object(c)
+            assert type(c) is t
+            assert c.real == 2
+            assert c.imag == 3
+            assert type(complex(c)) is complex
+            assert complex(c) == 2 + 3j
+            assert c == 2 + 3j
+            assert t(2) == 2
+            assert t(2.4) == 2.4
+
+            def assert_op_same(f):
+                assert f(t(2, 3)) == f(2 + 3j)
+
+            assert_op_same(abs)
+            assert_op_same(repr)
+            assert_op_same(format)
+            assert_op_same(bool)
+            assert_op_same(hash)
+            assert_op_same(operator.pos)
+            assert_op_same(operator.neg)
+            assert_op_same(lambda x: x + (1 + 2j))
+            assert_op_same(lambda x: x - (1 + 2j))
+            assert_op_same(lambda x: x * (1 + 2j))
+            assert_op_same(lambda x: x / (1 + 2j))
+            assert_op_same(lambda x: x ** (1 + 2j))
+            assert_op_same(lambda x: x.conjugate())