BTree subclasses can define max_bucket_size or max_btree_size to

  control maximum sizes for bucket and tree nodes.

BTrees/BTreeModuleTemplate.c

@@ -53,8 +53,8 @@
  */
 #define MODULE_NAME "BTrees." MOD_NAME_PREFIX "BTree."
 
-static PyObject *sort_str, *reverse_str, *__setstate___str,
-    *_bucket_type_str;
+static PyObject *sort_str, *reverse_str, *__setstate___str;
+static PyObject *_bucket_type_str, *max_btree_size_str, *max_bucket_size_str;
 static PyObject *ConflictError = NULL;
 
 static void PyVar_Assign(PyObject **v, PyObject *e) { Py_XDECREF(*v); *v=e;}
@@ -63,8 +63,6 @@ static void PyVar_Assign(PyObject **v, PyObject *e) { Py_XDECREF(*v); *v=e;}
 #define OBJECT(O) ((PyObject*)(O))
 
 #define MIN_BUCKET_ALLOC 16
-#define MAX_BTREE_SIZE(B) DEFAULT_MAX_BTREE_SIZE
-#define MAX_BUCKET_SIZE(B) DEFAULT_MAX_BUCKET_SIZE
 
 #define SameType_Check(O1, O2) (Py_TYPE((O1))==Py_TYPE((O2)))
 
@@ -223,6 +221,8 @@ typedef struct BTree_s {
    * data[len].key is positive infinity.
    */
   BTreeItem *data;
+  long max_btree_size;
+  long max_bucket_size;
 } BTree;
 
 static PyTypeObject BTreeType;
@@ -542,6 +542,13 @@ module_init(void)
     if (!_bucket_type_str)
         return NULL;
 
+    max_btree_size_str = INTERN("max_btree_size");
+    if (! max_btree_size_str)
+        return NULL;
+    max_bucket_size_str = INTERN("max_bucket_size");
+    if (! max_bucket_size_str)
+        return NULL;
+
     /* Grab the ConflictError class */
     interfaces = PyImport_ImportModule("BTrees.Interfaces");
     if (interfaces != NULL)

BTrees/BTreeTemplate.c

@@ -15,6 +15,50 @@
 
 #define BTREETEMPLATE_C "$Id$\n"
 
+static long
+_get_max_size(BTree *self, PyObject *name, long default_max)
+{
+  PyObject *size;
+  long isize;
+
+  size = PyObject_GetAttr(OBJECT(self->ob_type), name);
+  if (size == NULL) {
+    PyErr_Clear();
+    return default_max;
+  }
+  isize = PyInt_AsLong(size);
+  Py_DECREF(size);
+  if (isize <= 0 && ! PyErr_Occurred()) {
+    PyErr_SetString(PyExc_ValueError,
+                    "non-positive max size in BTree subclass");
+    return -1;
+  }
+
+  return isize;
+}
+
+static int
+_max_btree_size(BTree *self)
+{
+  long isize;
+
+  if (self->max_btree_size > 0) return self->max_btree_size;
+  isize = _get_max_size(self, max_btree_size_str, DEFAULT_MAX_BTREE_SIZE);
+  self->max_btree_size = isize;
+  return isize;
+}
+
+static int
+_max_bucket_size(BTree *self)
+{
+  long isize;
+
+  if (self->max_bucket_size > 0) return self->max_bucket_size;
+  isize = _get_max_size(self, max_bucket_size_str, DEFAULT_MAX_BUCKET_SIZE);
+  self->max_bucket_size = isize;
+  return isize;
+}
+
 /* Sanity-check a BTree.  This is a private helper for BTree_check.  Return:
  *      -1         Error.  If it's an internal inconsistency in the BTree,
  *                 AssertionError is set.
@@ -410,6 +454,9 @@ BTree_grow(BTree *self, int index, int noval)
 
     if (self->len)
     {
+        long max_size = _max_btree_size(self);
+        if (max_size < 0) return -1;
+
         d = self->data + index;
         v = d->child;
         /* Create a new object of the same type as the target value */
@@ -459,7 +506,7 @@ BTree_grow(BTree *self, int index, int noval)
         d->child = e;
         self->len++;
 
-        if (self->len >= MAX_BTREE_SIZE(self) * 2)    /* the root is huge */
+        if (self->len >= max_size * 2)    /* the root is huge */
             return BTree_split_root(self, noval);
     }
     else
@@ -727,11 +774,16 @@ _BTree_set(BTree *self, PyObject *keyarg, PyObject *value,
         int toobig;
 
         assert(status == 1);    /* can be 2 only on deletes */
-        if (SameType_Check(self, d->child))
-            toobig = childlength > MAX_BTREE_SIZE(d->child);
-        else
-            toobig = childlength > MAX_BUCKET_SIZE(d->child);
-
+        if (SameType_Check(self, d->child)) {
+            long max_size = _max_btree_size(self);
+            if (max_size < 0) return -1;
+            toobig = childlength > max_size;
+        }
+        else {
+            long max_size = _max_bucket_size(self);
+            if (max_size < 0) return -1;
+            toobig = childlength > max_size;
+        }
         if (toobig) {
             if (BTree_grow(self, min, noval) < 0)
                 goto Error;
@@ -2178,6 +2230,9 @@ BTree_init(PyObject *self, PyObject *args, PyObject *kwds)
 {
     PyObject *v = NULL;
 
+    BTREE(self)->max_bucket_size = 0;
+    BTREE(self)->max_btree_size = 0;
+
     if (!PyArg_ParseTuple(args, "|O:" MOD_NAME_PREFIX "BTree", &v))
         return -1;
 

BTrees/tests/test_btreesubclass.py

@@ -17,7 +17,14 @@ class B(OOBucket):
     pass
 
 class T(OOBTree):
+    __slots__ = ()
     _bucket_type = B
+    max_bucket_size = 2
+    max_btree_size = 3
+
+class S(T):
+    __slots__ = ()
+
 
 import unittest
 
@@ -27,18 +34,23 @@ class SubclassTest(unittest.TestCase):
         # test that a subclass that defines _bucket_type gets buckets
         # of that type
         t = T()
+        t[0] = 0
+        self.assertTrue(t._firstbucket.__class__ is B)
 
-        # There's no good way to get a bucket at the moment.
-        # __getstate__() is as good as it gets, but the default
-        # getstate explicitly includes the pickle of the bucket
-        # for small trees, so we have to be clever :-(
-
-        # make sure there is more than one bucket in the tree
-        for i in range(1000):
+    def testCustomNodeSizes(self):
+        # We overrise btree and bucket split sizes in BTree subclasses.
+        t = S()
+        for i in range(8):
             t[i] = i
-
-        state = t.__getstate__()
-        self.assertTrue(state[0][0].__class__ is B)
+        state = t.__getstate__()[0]
+        self.assertEqual(len(state), 5)
+        sub = state[0]
+        self.assertEqual(sub.__class__, S)
+        sub = sub.__getstate__()[0]
+        self.assertEqual(len(sub), 5)
+        sub = sub[0]
+        self.assertEqual(sub.__class__, B)
+        self.assertEqual(len(sub), 1)
 
 def test_suite():
     return unittest.makeSuite(SubclassTest)

CHANGES.rst

 ``BTrees`` Changelog
 ====================
 
+- BTree subclasses can define max_bucket_size or max_btree_size to
+  control maximum sizes for bucket and tree nodes.
 
 4.0.9 (unreleased)
 ------------------