gh-111495: Add tests for PyComplex C API

skirpichev · skirpichev · commit 637b759a6596 · 2023-11-01T08:16:44.000+03:00
diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -3,11 +3,14 @@
 from test import support
 from test.test_grammar import (VALID_UNDERSCORE_LITERALS,
                                INVALID_UNDERSCORE_LITERALS)
+from test.support import import_helper
 
 from random import random
 from math import atan2, isnan, copysign
 import operator
 
+_testcapi = import_helper.import_module('_testcapi')
+
 INF = float("inf")
 NAN = float("nan")
 # These tests ensure that complex math does the right thing
@@ -20,6 +23,8 @@
     (1, 0+0j),
 )
 
+NULL = None
+
 class ComplexTest(unittest.TestCase):
 
     def assertAlmostEqual(self, a, b):
@@ -72,6 +77,15 @@ def assertFloatsAreIdentical(self, x, y):
                 msg += ': zeros have different signs'
         self.fail(msg.format(x, y))
 
+    def assertComplexesAreIdentical(self, x, y):
+        """assert that complex numbers x and y are identical
+
+        I.e. they have identical real and imag components.
+
+        """
+        (self.assertFloatsAreIdentical(x.real, y.real)
+         and self.assertFloatsAreIdentical(x.imag, y.imag))
+
     def assertClose(self, x, y, eps=1e-9):
         """Return true iff complexes x and y "are close"."""
         self.assertCloseAbs(x.real, y.real, eps)
@@ -792,5 +806,94 @@ def test_format(self):
         self.assertEqual(format(complex(INF, -1), 'F'), 'INF-1.000000j')
 
 
+class CAPIComplexTest(ComplexTest):
+    def test_check(self):
+        # Test PyComplex_Check()
+        from test.test_capi.test_getargs import Complex, ComplexSubclass
+        check = _testcapi.complex_check
+
+        self.assertTrue(check(1+2j))
+        self.assertTrue(check(ComplexSubclass(1+2j)))
+        self.assertFalse(check(Complex()))
+        self.assertFalse(check(3))
+        self.assertFalse(check([]))
+        self.assertFalse(check(object()))
+
+        # CRASHES check(NULL)
+
+    def test_checkexact(self):
+        # PyComplex_CheckExact()
+        from test.test_capi.test_getargs import Complex, ComplexSubclass
+        checkexact = _testcapi.complex_checkexact
+
+        self.assertTrue(checkexact(1+2j))
+        self.assertFalse(checkexact(ComplexSubclass(1+2j)))
+        self.assertFalse(checkexact(Complex()))
+        self.assertFalse(checkexact(3))
+        self.assertFalse(checkexact([]))
+        self.assertFalse(checkexact(object()))
+
+        # CRASHES checkexact(NULL)
+
+    def test_fromccomplex(self):
+        # Test PyComplex_FromCComplex()
+        from test.test_capi.test_getargs import ComplexSubclass
+        fromccomplex = _testcapi.complex_fromccomplex
+
+        self.assertComplexesAreIdentical(fromccomplex(1+2j), 1.0+2.0j)
+
+    def test_fromdoubles(self):
+        # Test PyComplex_FromDoubles()
+        fromdoubles = _testcapi.complex_fromdoubles
+
+        self.assertComplexesAreIdentical(fromdoubles(1.0, 2.0), 1.0+2.0j)
+
+    def test_realasdouble(self):
+        # Test PyComplex_RealAsDouble()
+        from test.test_capi.test_getargs import BadComplex, Complex
+        realasdouble = _testcapi.complex_realasdouble
+
+        self.assertFloatsAreIdentical(realasdouble(1+2j), 1.0)
+        self.assertFloatsAreIdentical(realasdouble(1), 1.0)
+        self.assertFloatsAreIdentical(realasdouble(-1), -1.0)
+        # Function doesn't support classes with __complex__ dunder, see #109598
+        #self.assertFloatsAreIdentical(realasdouble(Complex()), 4.25)
+
+        self.assertRaises(TypeError, realasdouble, BadComplex())
+
+        # CRASHES realasdouble(NULL)
+
+    def test_imagasdouble(self):
+        # Test PyComplex_ImagAsDouble()
+        from test.test_capi.test_getargs import BadComplex, Complex
+        imagasdouble = _testcapi.complex_imagasdouble
+
+        self.assertFloatsAreIdentical(imagasdouble(1+2j), 2.0)
+        self.assertFloatsAreIdentical(imagasdouble(1), 0.0)
+        self.assertFloatsAreIdentical(imagasdouble(-1), 0.0)
+        # Function doesn't support classes with __complex__ dunder, see #109598
+        #self.assertFloatsAreIdentical(imagasdouble(Complex()), 0.5)
+
+        # Function returns 0.0 anyway, see #109598
+        #self.assertRaises(TypeError, imagasdouble, BadComplex())
+        self.assertFloatsAreIdentical(imagasdouble(BadComplex()), 0.0)
+
+        # CRASHES imagasdouble(NULL)
+
+    def test_asccomplex(self):
+        # Test PyComplex_AsCComplex()
+        from test.test_capi.test_getargs import BadComplex, Complex
+        asccomplex = _testcapi.complex_asccomplex
+
+        self.assertComplexesAreIdentical(asccomplex(1+2j), 1.0+2.0j)
+        self.assertComplexesAreIdentical(asccomplex(1), 1.0+0.0j)
+        self.assertComplexesAreIdentical(asccomplex(-1), -1.0+0.0j)
+        self.assertComplexesAreIdentical(asccomplex(Complex()), 4.25+0.5j)
+
+        self.assertRaises(TypeError, asccomplex, BadComplex())
+
+        # CRASHES asccomplex(NULL)
+
+
 if __name__ == "__main__":
     unittest.main()
diff --git a/Modules/_testcapi/complex.c b/Modules/_testcapi/complex.c
@@ -1,7 +1,104 @@
 #include "parts.h"
 #include "util.h"
 
+
+static PyObject *
+complex_check(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    NULLABLE(obj);
+    return PyLong_FromLong(PyComplex_Check(obj));
+}
+
+static PyObject *
+complex_checkexact(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    NULLABLE(obj);
+    return PyLong_FromLong(PyComplex_CheckExact(obj));
+}
+
+static PyObject *
+complex_fromccomplex(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    Py_complex complex = ((PyComplexObject*)obj)->cval;
+
+    return PyComplex_FromCComplex(complex);
+}
+
+static PyObject *
+complex_fromdoubles(PyObject *Py_UNUSED(module), PyObject *const *args,
+                    Py_ssize_t nargs)
+{
+    double real, imag;
+
+    assert(nargs == 2);
+
+    real = PyFloat_AsDouble(args[0]);
+    if (real == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    imag = PyFloat_AsDouble(args[1]);
+    if (imag == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyComplex_FromDoubles(real, imag);
+}
+
+static PyObject *
+complex_realasdouble(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    double real;
+
+    NULLABLE(obj);
+    real = PyComplex_RealAsDouble(obj);
+
+    if (real == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyFloat_FromDouble(real);
+}
+
+static PyObject *
+complex_imagasdouble(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    double imag;
+
+    NULLABLE(obj);
+    imag = PyComplex_ImagAsDouble(obj);
+
+    if (imag == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyFloat_FromDouble(imag);
+}
+
+static PyObject *
+complex_asccomplex(PyObject *Py_UNUSED(module), PyObject *obj)
+{
+    Py_complex complex;
+
+    NULLABLE(obj);
+    complex = PyComplex_AsCComplex(obj);
+
+    if (complex.real == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyComplex_FromCComplex(complex);
+}
+
+
 static PyMethodDef test_methods[] = {
+    {"complex_check", complex_check, METH_O},
+    {"complex_checkexact", complex_checkexact, METH_O},
+    {"complex_fromccomplex", complex_fromccomplex, METH_O},
+    {"complex_fromdoubles", _PyCFunction_CAST(complex_fromdoubles), METH_FASTCALL},
+    {"complex_realasdouble", complex_realasdouble, METH_O},
+    {"complex_imagasdouble", complex_imagasdouble, METH_O},
+    {"complex_asccomplex", complex_asccomplex, METH_O},
     {NULL},
 };
 
