Merge the alienoid-btrees_setdefault branch.

This adds .setdefault() methods to BTrees and Buckets.

NEWS.txt

@@ -183,6 +183,13 @@ FileStorage
 BTrees
 ------
 
+- (3.5a9) BTrees and Buckets now implement the ``setdefault()`` method.
+  This is exactly list Python's ``setdefault()`` method for dictionaries,
+  except that both arguments are required (and Python is likely to change
+  to require both arguments too -- defaulting the ``default`` argument to
+  ``None`` has no viable use cases).  Thanks to Ruslan Spivak for
+  contributing code, tests, and documentation.
+
 - (3.5a5) Collector 1843.  When a non-integer was passed to a method like
   ``keys()`` of a Bucket or Set with integer keys, an internal error code
   was overlooked, leading to everything from "delayed errors" to segfaults.

src/BTrees/BTreeTemplate.c

@@ -1700,6 +1700,36 @@ BTree_has_key(BTree *self, PyObject *key)
     return _BTree_get(self, key, 1);
 }
 
+static PyObject *
+BTree_setdefault(BTree *self, PyObject *args)
+{
+    PyObject *key;
+    PyObject *failobj; /* default */
+    PyObject *value;   /* return value */
+
+    if (! PyArg_UnpackTuple(args, "setdefault", 2, 2, &key, &failobj))
+    	return NULL;
+
+    value = _BTree_get(self, key, 0);
+    if (value != 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;
+    PyErr_Clear();
+
+    /* Associate `key` with `failobj` in the tree, and return `failobj`. */
+    value = failobj;
+    if (_BTree_set(self, key, failobj, 0, 0) < 0)
+        value = NULL;
+    Py_XINCREF(value);
+    return value;
+}
+
+
 /* Search BTree self for key.  This is the sq_contains slot of the
  * PySequenceMethods.
  *
@@ -1838,6 +1868,11 @@ static struct PyMethodDef BTree_methods[] = {
      "get(key[, default=None]) -> Value for key or default\n\n"
      "Return the value or the default if the key is not found."},
 
+    {"setdefault", (PyCFunction) BTree_setdefault, METH_VARARGS,
+     "D.setdefault(k, d) -> D.get(k, d), also set D[k]=d if k not in D.\n\n"
+     "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."},
+
     {"maxKey", (PyCFunction) BTree_maxKey,	METH_VARARGS,
      "maxKey([max]) -> key\n\n"
      "Return the largest key in the BTree.  If max is specified, return\n"

src/BTrees/BucketTemplate.c

@@ -1257,6 +1257,35 @@ bucket_has_key(Bucket *self, PyObject *key)
     return _bucket_get(self, key, 1);
 }
 
+static PyObject *
+bucket_setdefault(Bucket *self, PyObject *args)
+{
+    PyObject *key;
+    PyObject *failobj; /* default */
+    PyObject *value;   /* return value */
+    int dummy_changed; /* in order to call _bucket_set */
+
+    if (! PyArg_UnpackTuple(args, "setdefault", 2, 2, &key, &failobj))
+    	return NULL;
+
+    value = _bucket_get(self, key, 0);
+    if (value != 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;
+    PyErr_Clear();
+
+    /* Associate `key` with `failobj` in the bucket, and return `failobj`. */
+    value = failobj;
+    if (_bucket_set(self, key, failobj, 0, 0, &dummy_changed) < 0)
+        value = NULL;
+    Py_XINCREF(value);
+    return value;
+}
 
 /* Search bucket self for key.  This is the sq_contains slot of the
  * PySequenceMethods.
@@ -1481,6 +1510,11 @@ static struct PyMethodDef Bucket_methods[] = {
      "get(key[,default]) -- Look up a value\n\n"
      "Return the default (or None) if the key is not found."},
 
+    {"setdefault", (PyCFunction) bucket_setdefault, METH_VARARGS,
+     "D.setdefault(k, d) -> D.get(k, d), also set D[k]=d if k not in D.\n\n"
+     "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."},
+
     {"iterkeys", (PyCFunction) Bucket_iterkeys,  METH_KEYWORDS,
      "B.iterkeys([min[,max]]) -> an iterator over the keys of B"},
 

src/BTrees/Interfaces.py

@@ -214,6 +214,18 @@ class IDictionaryIsh(IMinimalDictionary):
         integer values, the normalization is division.
         """
 
+    def setdefault(key, d):
+        """D.setdefault(k, d) -> D.get(k, d), also set D[k]=d if k not in D.
+
+        Return the value like get() except that if key is missing, d is both
+        returned and inserted into the dictionary as the value of k.
+
+        Note that, unlike as for Python's dict.setdefault(), d is not
+        optional.  Python defaults d to None, but that doesn't make sense
+        for mappings that can't have None as a value (for example, an
+        IIBTree can have only integers as values).
+        """
+
 class IBTree(IDictionaryIsh):
 
     def insert(key, value):

src/BTrees/tests/testBTrees.py

@@ -611,6 +611,23 @@ class MappingBase(Base):
                                        excludemax=True)),
                              f([1]))
 
+    def testSetdefault(self):
+        t = self.t
+
+        self.assertEqual(t.setdefault(1, 2), 2)
+        # That should also have associated 1 with 2 in the tree.
+        self.assert_(1 in t)
+        self.assertEqual(t[1], 2)
+        # And trying to change it again should have no effect.
+        self.assertEqual(t.setdefault(1, 666), 2)
+        self.assertEqual(t[1], 2)
+
+        # Not enough arguments.
+        self.assertRaises(TypeError, t.setdefault)
+        self.assertRaises(TypeError, t.setdefault, 1)
+        # Too many arguments.
+        self.assertRaises(TypeError, t.setdefault, 1, 2, 3)
+
 class NormalSetTests(Base):
     """ Test common to all set types """