Do the same thing to complex that I did to str: the rich comparison

function returns NotImplemented when comparing objects whose
tp_richcompare slot is not itself.

diff --git a/Lib/test/test_descr.py b/Lib/test/test_descr.py
index ed3cea44b896ee8415d54228d5fe21d5fbf99463..42e13846a71a0b6fbc931db3ab73ad9380da3896 100644 (file)
--- a/Lib/test/test_descr.py
+++ b/Lib/test/test_descr.py
@@ -1863,6 +1863,21 @@ def classic_comparisons():
 def rich_comparisons():
     if verbose:
         print "Testing rich comparisons..."
+    class Z(complex):
+        pass
+    z = Z(1)
+    verify(z == 1+0j)
+    verify(1+0j == z)
+    class ZZ(complex):
+        def __eq__(self, other):
+            try:
+                return abs(self - other) <= 1e-6
+            except:
+                return NotImplemented
+    zz = ZZ(1.0000003)
+    verify(zz == 1+0j)
+    verify(1+0j == zz)
+            
     class classic:
         pass
     for base in (classic, int, object, list):

diff --git a/Objects/complexobject.c b/Objects/complexobject.c
index 191dcba73eec33a6d669f8c3b8ea255b67cf19f0..32f2b2459fe33e845f164682fdd681632a1e9d33 100644 (file)
--- a/Objects/complexobject.c
+++ b/Objects/complexobject.c
@@ -553,12 +553,6 @@ complex_richcompare(PyObject *v, PyObject *w, int op)
        Py_complex i, j;
        PyObject *res;
 
-       if (op != Py_EQ && op != Py_NE) {
-               PyErr_SetString(PyExc_TypeError,
-                       "cannot compare complex numbers using <, <=, >, >=");
-               return NULL;
-       }
-
        c = PyNumber_CoerceEx(&v, &w);
        if (c < 0)
                return NULL;
@@ -566,7 +560,10 @@ complex_richcompare(PyObject *v, PyObject *w, int op)
                Py_INCREF(Py_NotImplemented);
                return Py_NotImplemented;
        }
-       if (!PyComplex_Check(v) || !PyComplex_Check(w)) {
+       /* May sure both arguments use complex comparison.
+          This implies PyComplex_Check(a) && PyComplex_Check(b). */
+       if (v->ob_type->tp_richcompare != complex_richcompare ||
+           w->ob_type->tp_richcompare != complex_richcompare) {
                Py_DECREF(v);
                Py_DECREF(w);
                Py_INCREF(Py_NotImplemented);
@@ -578,6 +575,12 @@ complex_richcompare(PyObject *v, PyObject *w, int op)
        Py_DECREF(v);
        Py_DECREF(w);
 
+       if (op != Py_EQ && op != Py_NE) {
+               PyErr_SetString(PyExc_TypeError,
+                       "cannot compare complex numbers using <, <=, >, >=");
+               return NULL;
+       }
+
        if ((i.real == j.real && i.imag == j.imag) == (op == Py_EQ))
                res = Py_True;
        else