Patch #1675423: PyComplex_AsCComplex() now tries to convert an object

to complex using its __complex__() method before falling back to the
__float__() method. Therefore, the functions in the cmath module now
can operate on objects that define a __complex__() method.
 (backport)

1 files changed, 48 insertions, 1 deletions

diff --git a/Objects/complexobject.c b/Objects/complexobject.c
index 54ad075..a8a2188 100644
--- a/Objects/complexobject.c
+++ b/Objects/complexobject.c
@@ -252,12 +252,59 @@ Py_complex
 PyComplex_AsCComplex(PyObject *op)
 {
 	Py_complex cv;
+	PyObject *newop = NULL;
+	static PyObject *complex_str = NULL;
+
+	assert(op);
+	/* If op is already of type PyComplex_Type, return its value */
 	if (PyComplex_Check(op)) {
 		return ((PyComplexObject *)op)->cval;
 	}
+	/* If not, use op's __complex__  method, if it exists */
+	
+	/* return -1 on failure */
+	cv.real = -1.;
+	cv.imag = 0.;
+	
+	if (PyInstance_Check(op)) {
+		/* this can go away in python 3000 */
+		if (PyObject_HasAttrString(op, "__complex__")) {
+			newop = PyObject_CallMethod(op, "__complex__", NULL);
+			if (!newop)
+				return cv;
+		}
+		/* else try __float__ */
+	} else {
+		PyObject *complexfunc;
+		if (!complex_str) {
+			if (!(complex_str = PyString_FromString("__complex__")))
+				return cv;
+		}
+		complexfunc = _PyType_Lookup(op->ob_type, complex_str);
+		/* complexfunc is a borrowed reference */
+		if (complexfunc) {
+			newop = PyObject_CallFunctionObjArgs(complexfunc, op, NULL);
+			if (!newop)
+				return cv;
+		}
+	}
+
+	if (newop) {
+		if (!PyComplex_Check(newop)) {
+			PyErr_SetString(PyExc_TypeError,
+				"__complex__ should return a complex object");
+			Py_DECREF(newop);
+			return cv;
+		}
+		cv = ((PyComplexObject *)newop)->cval;
+		Py_DECREF(newop);
+		return cv;
+	}
+	/* If neither of the above works, interpret op as a float giving the
+	   real part of the result, and fill in the imaginary part as 0. */
 	else {
+		/* PyFloat_AsDouble will return -1 on failure */
 		cv.real = PyFloat_AsDouble(op);
-		cv.imag = 0.;
 		return cv;
 	}
 }