Rev 4688: Bootstrap a richcompare function using as_tuples. in http://bazaar.launchpad.net/~jameinel/bzr/2.1-memory-consumption

[John Arbash Meinel]

At http://bazaar.launchpad.net/~jameinel/bzr/2.1-memory-consumption

------------------------------------------------------------
revno: 4688
revision-id: john at arbash-meinel.com-20090909153943-ryi4hotash4l8n46
parent: john at arbash-meinel.com-20090909152038-j5ggmryfyznohcas
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 2.1-memory-consumption
timestamp: Wed 2009-09-09 10:39:43 -0500
message:
  Bootstrap a richcompare function using as_tuples.
  This allows 'sorted()' etc to function properly.
-------------- next part --------------
=== modified file 'bzrlib/_keys_type_c.c'
--- a/bzrlib/_keys_type_c.c	2009-09-09 15:20:38 +0000
+++ b/bzrlib/_keys_type_c.c	2009-09-09 15:39:43 +0000
@@ -42,6 +42,8 @@
 extern PyTypeObject KeysType;
 static PyObject *Keys_item(Keys *self, Py_ssize_t offset);
 
+#define Keys_CheckExact(op) (Py_TYPE(op) == &KeysType)
+
 static void
 Keys_dealloc(Keys *keys)
 {
@@ -156,6 +158,41 @@
 }
 
 static PyObject *
+Keys_richcompare(PyObject *v, PyObject *w, int op)
+{
+    PyObject *vt, *wt;
+    PyObject *vt_to_decref = NULL, *wt_to_decref = NULL;
+    PyObject *result;
+    
+    if (Keys_CheckExact(v)) {
+        vt = Keys_as_tuples((Keys *)v);
+        vt_to_decref = vt;
+    } else if (PyTuple_Check(v)) {
+        vt = v;
+        vt_to_decref = NULL;
+    } else {
+		Py_INCREF(Py_NotImplemented);
+		return Py_NotImplemented;
+    }
+    if (Keys_CheckExact(w)) {
+        wt = Keys_as_tuples((Keys *)w);
+        wt_to_decref = wt;
+    } else if (PyTuple_Check(w)) {
+        wt = w;
+        wt_to_decref = NULL;
+    } else {
+        Py_XDECREF(vt_to_decref);
+		Py_INCREF(Py_NotImplemented);
+		return Py_NotImplemented;
+    }
+    /* Now we have 2 tuples to compare, do it */
+    result = PyTuple_Type.tp_richcompare(vt, wt, op);
+    Py_XDECREF(vt_to_decref);
+    Py_XDECREF(wt_to_decref);
+    return result;
+}
+
+static PyObject *
 Keys_repr(Keys *self)
 {
     PyObject *as_tpl;
@@ -166,7 +203,6 @@
         return NULL;
     }
     result = PyObject_Repr(as_tpl);
-Done:
     Py_DECREF(as_tpl);
     return result;
 }
@@ -244,7 +280,7 @@
     0,                                           /* tp_setattr */
     0,                                           /* tp_compare */
     // TODO: implement repr() and possibly str()
-    Keys_repr,                                           /* tp_repr */
+    (reprfunc)Keys_repr,                         /* tp_repr */
     0,                                           /* tp_as_number */
     &Keys_as_sequence,                           /* tp_as_sequence */
     0,                                           /* tp_as_mapping */
@@ -263,7 +299,7 @@
     0,                                           /* tp_clear */
     // TODO: implement richcompare, we should probably be able to compare vs an
     //       tuple, as well as versus another Keys object.
-    0,                                           /* tp_richcompare */
+    Keys_richcompare,                            /* tp_richcompare */
     0,                                           /* tp_weaklistoffset */
     // We could implement this as returning tuples of keys...
     0,                                           /* tp_iter */

=== modified file 'bzrlib/tests/test__keys_type.py'
--- a/bzrlib/tests/test__keys_type.py	2009-09-09 15:20:38 +0000
+++ b/bzrlib/tests/test__keys_type.py	2009-09-09 15:39:43 +0000
@@ -120,3 +120,30 @@
     def test_repr(self):
         k = self.module.Keys(2, 'foo', 'bar', 'baz', 'bing')
         self.assertEqual("(('foo', 'bar'), ('baz', 'bing'))", repr(k))
+
+    def test_compare(self):
+        k1 = self.module.Keys(2, 'foo', 'bar')
+        k2 = self.module.Keys(2, 'baz', 'bing')
+        k3 = self.module.Keys(2, 'foo', 'zzz')
+        k4 = self.module.Keys(2, 'foo', 'bar')
+        # Comparison should be done on the keys themselves, and not based on
+        # object id, etc.
+        self.assertTrue(k1 == k1)
+        self.assertTrue(k1 == k4)
+        self.assertTrue(k2 < k1)
+        self.assertTrue(k2 < k4)
+        self.assertTrue(k3 > k1)
+        self.assertTrue(k3 > k4)
+        # We should also be able to compare against raw tuples
+        self.assertTrue(k1 == (('foo', 'bar'),))
+
+    def test_sorted(self):
+        k1 = self.module.Keys(2, 'foo', 'bar', 'baz', 'bing', 'foo', 'zzz')
+        self.assertEqual([('baz', 'bing'), ('foo', 'bar'), ('foo', 'zzz')],
+                         sorted(k1))
+
+        k1 = self.module.Keys(2, 'foo', 'bar')
+        k2 = self.module.Keys(2, 'baz', 'bing')
+        k3 = self.module.Keys(2, 'foo', 'zzz')
+        self.assertEqual([(('baz', 'bing'),), (('foo', 'bar'),),
+                          (('foo', 'zzz'),)], sorted([k1, k2, k3]))