[Zodb-checkins] SVN: ZODB/branches/alienoid-btrees_pop/src/BTrees/ Added pop method implementation.

[Ruslan Spivak]

Log message for revision 38165:
  Added pop method implementation.

Changed:
  U   ZODB/branches/alienoid-btrees_pop/src/BTrees/BTreeTemplate.c
  U   ZODB/branches/alienoid-btrees_pop/src/BTrees/BucketTemplate.c
  U   ZODB/branches/alienoid-btrees_pop/src/BTrees/tests/testBTrees.py

-=-
Modified: ZODB/branches/alienoid-btrees_pop/src/BTrees/BTreeTemplate.c
===================================================================
--- ZODB/branches/alienoid-btrees_pop/src/BTrees/BTreeTemplate.c	2005-08-30 20:38:50 UTC (rev 38164)
+++ ZODB/branches/alienoid-btrees_pop/src/BTrees/BTreeTemplate.c	2005-08-30 20:42:14 UTC (rev 38165)
@@ -1729,7 +1729,60 @@
     return value;
 }
 
+/* forward declaration */
+static int
+BTree_length_or_nonzero(BTree *self, int nonzero);
 
+static PyObject *
+BTree_pop(BTree *self, PyObject *args)
+{
+    PyObject *key;
+    PyObject *failobj = NULL; /* default */
+    PyObject *value;          /* return value */
+
+    if (! PyArg_UnpackTuple(args, "pop", 1, 2, &key, &failobj))
+    	return NULL;
+
+    value = _BTree_get(self, key, 0);
+    if (value != NULL) {
+        if (_BTree_set(self, key, 0, 0, 0) < 0) {
+            Py_DECREF(value);
+            return NULL;;
+        }
+        return value;
+    }
+
+    /* The key isn't in the tree.  If that's not due to a KeyError exception,
+     * pass back the unexpected exception.
+     */
+    if (! PyErr_ExceptionMatches(PyExc_KeyError))
+        return NULL;
+
+    int nonzero = 1;
+    int res = BTree_length_or_nonzero(self, nonzero);
+    if (res < 0) {
+        return NULL;
+    } else if (!res) {
+        /* btree is empty */
+        if (failobj == NULL) {
+            PyErr_SetString(PyExc_KeyError, "pop(): dictionary is empty");
+            return NULL;
+        }
+        PyErr_Clear();
+        Py_INCREF(failobj);
+        return failobj;
+    }
+
+    /* btree is not empty */
+    if (failobj == NULL)
+        return NULL;
+
+    PyErr_Clear();
+    Py_INCREF(failobj);
+    return failobj;
+}
+
+
 /* Search BTree self for key.  This is the sq_contains slot of the
  * PySequenceMethods.
  *
@@ -1873,6 +1926,10 @@
      "Return the value like get() except that if key is missing, d is both\n"
      "returned and inserted into the BTree as the value of k."},
 
+    {"pop", (PyCFunction) BTree_pop, METH_VARARGS,
+     "D.pop(k[,d]) -> v, remove specified key and return the corresponding value.\n\n"
+     "If key is not found, d is returned if given, otherwise KeyError is raised."},
+
     {"maxKey", (PyCFunction) BTree_maxKey,	METH_VARARGS,
      "maxKey([max]) -> key\n\n"
      "Return the largest key in the BTree.  If max is specified, return\n"

Modified: ZODB/branches/alienoid-btrees_pop/src/BTrees/BucketTemplate.c
===================================================================
--- ZODB/branches/alienoid-btrees_pop/src/BTrees/BucketTemplate.c	2005-08-30 20:38:50 UTC (rev 38164)
+++ ZODB/branches/alienoid-btrees_pop/src/BTrees/BucketTemplate.c	2005-08-30 20:42:14 UTC (rev 38165)
@@ -1287,6 +1287,61 @@
     return value;
 }
 
+
+/* forward declaration */
+static int
+Bucket_length(Bucket *self);
+
+static PyObject *
+bucket_pop(Bucket *self, PyObject *args)
+{
+    PyObject *key;
+    PyObject *failobj = NULL; /* default */
+    PyObject *value;          /* return value */
+    int dummy_changed;        /* in order to call _bucket_set */
+
+    if (! PyArg_UnpackTuple(args, "pop", 1, 2, &key, &failobj))
+    	return NULL;
+
+    value = _bucket_get(self, key, 0);
+    if (value != NULL) {
+        if (_bucket_set(self, key, 0, 0, 0, &dummy_changed) < 0) {
+            Py_DECREF(value);
+            return NULL;;
+        }
+        return value;
+    }
+
+    /* The key isn't in the bucket.  If that's not due to a KeyError exception,
+     * pass back the unexpected exception.
+     */
+    if (! PyErr_ExceptionMatches(PyExc_KeyError))
+        return NULL;
+
+    int res = Bucket_length(self);
+    if (res < 0) {
+        return NULL;
+    } else if (!res) {
+        /* bucket is empty */
+        if (failobj == NULL) {
+            PyErr_SetString(PyExc_KeyError, "pop(): dictionary is empty");
+            return NULL;
+        }
+        PyErr_Clear();
+        Py_INCREF(failobj);
+        return failobj;
+    }
+
+    /* bucket is not empty */
+    if (failobj == NULL)
+        return NULL;
+
+    PyErr_Clear();
+    Py_INCREF(failobj);
+    return failobj;
+}
+
+
 /* Search bucket self for key.  This is the sq_contains slot of the
  * PySequenceMethods.
  *
@@ -1515,6 +1570,10 @@
      "Return the value like get() except that if key is missing, d is both\n"
      "returned and inserted into the bucket as the value of k."},
 
+    {"pop", (PyCFunction) bucket_pop, METH_VARARGS,
+     "D.pop(k[,d]) -> v, remove specified key and return the corresponding value.\n\n"
+     "If key is not found, d is returned if given, otherwise KeyError is raised."},
+
     {"iterkeys", (PyCFunction) Bucket_iterkeys,  METH_KEYWORDS,
      "B.iterkeys([min[,max]]) -> an iterator over the keys of B"},
 

Modified: ZODB/branches/alienoid-btrees_pop/src/BTrees/tests/testBTrees.py
===================================================================
--- ZODB/branches/alienoid-btrees_pop/src/BTrees/tests/testBTrees.py	2005-08-30 20:38:50 UTC (rev 38164)
+++ ZODB/branches/alienoid-btrees_pop/src/BTrees/tests/testBTrees.py	2005-08-30 20:42:14 UTC (rev 38165)
@@ -628,6 +628,21 @@
         # Too many arguments.
         self.assertRaises(TypeError, t.setdefault, 1, 2, 3)
 
+
+    def testPop(self):
+        t = self.t
+
+        # KeyError when container is empty
+        self.assertRaises(KeyError, t.pop, 1)
+
+        t[1] = 3
+        self.assertRaises(KeyError, t.pop, 5)
+        self.assertEqual(t.pop(1), 3)
+        self.assertEqual(len(t), 0)
+
+        self.assertEqual(t.pop(1, 3), 3)
+
+
 class NormalSetTests(Base):
     """ Test common to all set types """