python · encukou · Jan 15, 2024 · Sep 21, 2023 · Oct 10, 2023 · Oct 10, 2023
diff --git a/Doc/c-api/complex.rst b/Doc/c-api/complex.rst
@@ -117,11 +117,29 @@ Complex Numbers as Python Objects
 
    Return the real part of *op* as a C :c:expr:`double`.
 
+   If *op* is not a Python complex number object but has a
+   :meth:`~object.__complex__` method, this method will first be called to
+   convert *op* to a Python complex number object.  If :meth:`!__complex__` is
+   not defined then it falls back to call :c:func:`PyFloat_AsDouble` and
+   returns its result.  Upon failure, this method returns ``-1.0``, so one
+   should call :c:func:`PyErr_Occurred` to check for errors.
+
+   .. versionchanged:: 3.13
+      Use :meth:`~object.__complex__` if available.
 
 .. c:function:: double PyComplex_ImagAsDouble(PyObject *op)
 
    Return the imaginary part of *op* as a C :c:expr:`double`.
 
+   If *op* is not a Python complex number object but has a
+   :meth:`~object.__complex__` method, this method will first be called to
+   convert *op* to a Python complex number object.  If :meth:`!__complex__` is
+   not defined then it falls back to call :c:func:`PyFloat_AsDouble` and
+   returns ``0.0`` on success.  Upon failure, this method returns ``-1.0``, so
+   one should call :c:func:`PyErr_Occurred` to check for errors.
+
+   .. versionchanged:: 3.13
+      Use :meth:`~object.__complex__` if available.
 
 .. c:function:: Py_complex PyComplex_AsCComplex(PyObject *op)
 

diff --git a/Lib/test/test_complex.py b/Lib/test/test_complex.py
@@ -8,6 +8,11 @@
 from math import atan2, isnan, copysign
 import operator
 
+try:
+    import _testcapi
+except ImportError:
+    _testcapi = None
+
 INF = float("inf")
 NAN = float("nan")
 # These tests ensure that complex math does the right thing
@@ -791,6 +796,34 @@ def test_format(self):
         self.assertEqual(format(complex(INF, 1), 'F'), 'INF+1.000000j')
         self.assertEqual(format(complex(INF, -1), 'F'), 'INF-1.000000j')
 
+    @unittest.skipIf(_testcapi is None, 'needs _testcapi')
+    def test_PyComplex_RealAsDouble(self):
+        from test.test_capi.test_getargs import BadComplex, Complex
+
+        f = _testcapi.complex_real_as_double
+
+        self.assertFloatsAreIdentical(f(1+2j), 1.0)
+        self.assertFloatsAreIdentical(f(1), 1.0)
+        self.assertFloatsAreIdentical(f(-1), -1.0)
+        self.assertFloatsAreIdentical(f(Complex()), 4.25)
+
+        self.assertRaises(TypeError, f, None)
+        self.assertRaises(TypeError, f, BadComplex())
+
+    @unittest.skipIf(_testcapi is None, 'needs _testcapi')
+    def test_PyComplex_ImagAsDouble(self):
+        from test.test_capi.test_getargs import BadComplex, Complex
+
+        f = _testcapi.complex_imag_as_double
+
+        self.assertFloatsAreIdentical(f(1+2j), 2.0)
+        self.assertFloatsAreIdentical(f(1), 0.0)
+        self.assertFloatsAreIdentical(f(-1), 0.0)
+        self.assertFloatsAreIdentical(f(Complex()), 0.5)
+
+        self.assertRaises(TypeError, f, None)
+        self.assertRaises(TypeError, f, BadComplex())
+
 
 if __name__ == "__main__":
     unittest.main()
diff --git a/Misc/NEWS.d/next/Core and Builtins/2023-09-21-11-54-28.gh-issue-109598.CRidSy.rst b/Misc/NEWS.d/next/Core and Builtins/2023-09-21-11-54-28.gh-issue-109598.CRidSy.rst
@@ -0,0 +1,3 @@
+:c:func:`PyComplex_RealAsDouble`/:c:func:`PyComplex_ImagAsDouble` now tries to
+convert an object to a :class:`complex` instance using its ``__complex__()`` method
+before falling back to the ``__float__()`` method.  Patch by Sergey B Kirpichev.
diff --git a/Modules/Setup.stdlib.in b/Modules/Setup.stdlib.in
@@ -159,7 +159,7 @@
 @MODULE__XXTESTFUZZ_TRUE@_xxtestfuzz _xxtestfuzz/_xxtestfuzz.c _xxtestfuzz/fuzzer.c
 @MODULE__TESTBUFFER_TRUE@_testbuffer _testbuffer.c
 @MODULE__TESTINTERNALCAPI_TRUE@_testinternalcapi _testinternalcapi.c _testinternalcapi/test_lock.c _testinternalcapi/pytime.c
-@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/pyos.c _testcapi/immortal.c _testcapi/heaptype_relative.c _testcapi/gc.c
+@MODULE__TESTCAPI_TRUE@_testcapi _testcapimodule.c _testcapi/vectorcall.c _testcapi/vectorcall_limited.c _testcapi/heaptype.c _testcapi/abstract.c _testcapi/unicode.c _testcapi/dict.c _testcapi/getargs.c _testcapi/datetime.c _testcapi/docstring.c _testcapi/mem.c _testcapi/watchers.c _testcapi/long.c _testcapi/float.c _testcapi/structmember.c _testcapi/exceptions.c _testcapi/code.c _testcapi/buffer.c _testcapi/pyatomic.c _testcapi/pyos.c _testcapi/immortal.c _testcapi/heaptype_relative.c _testcapi/gc.c _testcapi/complex.c
 @MODULE__TESTCLINIC_TRUE@_testclinic _testclinic.c
 @MODULE__TESTCLINIC_LIMITED_TRUE@_testclinic_limited _testclinic_limited.c
 

diff --git a/Modules/_testcapi/clinic/complex.c.h b/Modules/_testcapi/clinic/complex.c.h
diff --git a/Modules/_testcapi/complex.c b/Modules/_testcapi/complex.c
@@ -0,0 +1,68 @@
+#include "parts.h"
+#include "clinic/complex.c.h"
+
+
+/*[clinic input]
+module _testcapi
+[clinic start generated code]*/
+/*[clinic end generated code: output=da39a3ee5e6b4b0d input=6361033e795369fc]*/
+
+/*[clinic input]
+_testcapi.complex_real_as_double
+
+    op: object
+    /
+
+Test PyComplex_RealAsDouble().
+[clinic start generated code]*/
+
+static PyObject *
+_testcapi_complex_real_as_double(PyObject *module, PyObject *op)
+/*[clinic end generated code: output=91427770a4583095 input=7ffd3ffd53495b3d]*/
+{
+    double real = PyComplex_RealAsDouble(op);
+
+    if (real == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyFloat_FromDouble(real);
+}
+
+/*[clinic input]
+_testcapi.complex_imag_as_double
+
+    op: object
+    /
+
+Test PyComplex_ImagAsDouble().
+[clinic start generated code]*/
+
+static PyObject *
+_testcapi_complex_imag_as_double(PyObject *module, PyObject *op)
+/*[clinic end generated code: output=f66f10a3d47beec4 input=e2b4b00761e141ea]*/
+{
+    double imag = PyComplex_ImagAsDouble(op);
+
+    if (imag == -1. && PyErr_Occurred()) {
+        return NULL;
+    }
+
+    return PyFloat_FromDouble(imag);
+}
+
+static PyMethodDef test_methods[] = {
+    _TESTCAPI_COMPLEX_REAL_AS_DOUBLE_METHODDEF
+    _TESTCAPI_COMPLEX_IMAG_AS_DOUBLE_METHODDEF
+    {NULL},
+};
+
+int
+_PyTestCapi_Init_Complex(PyObject *mod)
+{
+    if (PyModule_AddFunctions(mod, test_methods) < 0) {
+        return -1;
+    }
+
+    return 0;
+}
diff --git a/Modules/_testcapi/parts.h b/Modules/_testcapi/parts.h
@@ -26,6 +26,7 @@ int _PyTestCapi_Init_PyAtomic(PyObject *module);
 int _PyTestCapi_Init_PyOS(PyObject *module);
 int _PyTestCapi_Init_Immortal(PyObject *module);
 int _PyTestCapi_Init_GC(PyObject *mod);
+int _PyTestCapi_Init_Complex(PyObject *mod);
 
 int _PyTestCapi_Init_VectorcallLimited(PyObject *module);
 int _PyTestCapi_Init_HeaptypeRelative(PyObject *module);

diff --git a/Modules/_testcapimodule.c b/Modules/_testcapimodule.c
@@ -4024,6 +4024,9 @@ PyInit__testcapi(void)
     if (_PyTestCapi_Init_HeaptypeRelative(m) < 0) {
         return NULL;
     }
+    if (_PyTestCapi_Init_Complex(m) < 0) {
+        return NULL;
+    }
 
     PyState_AddModule(m, &_testcapimodule);
     return m;

diff --git a/Objects/complexobject.c b/Objects/complexobject.c
@@ -256,26 +256,51 @@ PyComplex_FromDoubles(double real, double imag)
     return PyComplex_FromCComplex(c);
 }
 
+static PyObject * try_complex_special_method(PyObject *);
+
 double
 PyComplex_RealAsDouble(PyObject *op)
 {
+    double real = -1.0;
+
     if (PyComplex_Check(op)) {
-        return ((PyComplexObject *)op)->cval.real;
+        real = ((PyComplexObject *)op)->cval.real;
     }
     else {
-        return PyFloat_AsDouble(op);
+        PyObject* newop = try_complex_special_method(op);
+        if (newop) {
+            real = ((PyComplexObject *)newop)->cval.real;
+            Py_DECREF(newop);
+        } else if (!PyErr_Occurred()) {
+            real = PyFloat_AsDouble(op);
+        }
     }
+
+    return real;
 }
 
 double
 PyComplex_ImagAsDouble(PyObject *op)
 {
+    double imag = -1.0;
+
     if (PyComplex_Check(op)) {
-        return ((PyComplexObject *)op)->cval.imag;
+        imag = ((PyComplexObject *)op)->cval.imag;
     }
     else {
-        return 0.0;
+        PyObject* newop = try_complex_special_method(op);
+        if (newop) {
+            imag = ((PyComplexObject *)newop)->cval.imag;
+            Py_DECREF(newop);
+        } else if (!PyErr_Occurred()) {
+            PyFloat_AsDouble(op);
+            if (!PyErr_Occurred()) {
+                imag = 0.0;
+            }
+        }
     }
+
+    return imag;
 }
 
 static PyObject *

diff --git a/PCbuild/_testcapi.vcxproj b/PCbuild/_testcapi.vcxproj
@@ -116,6 +116,7 @@
     <ClCompile Include="..\Modules\_testcapi\pyos.c" />
     <ClCompile Include="..\Modules\_testcapi\immortal.c" />
     <ClCompile Include="..\Modules\_testcapi\gc.c" />
+    <ClCompile Include="..\Modules\_testcapi\complex.c" />
   </ItemGroup>
   <ItemGroup>
     <ResourceCompile Include="..\PC\python_nt.rc" />

diff --git a/PCbuild/_testcapi.vcxproj.filters b/PCbuild/_testcapi.vcxproj.filters
@@ -75,6 +75,9 @@
     <ClCompile Include="..\Modules\_testcapi\gc.c">
       <Filter>Source Files</Filter>
     </ClCompile>
+    <ClCompile Include="..\Modules\_testcapi\complex.c">
+      <Filter>Source Files</Filter>
+    </ClCompile>
   </ItemGroup>
   <ItemGroup>
     <ResourceCompile Include="..\PC\python_nt.rc">