Rev 4681: Add a get_key() function, which then returns tuples for the given items. 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: 4681
revision-id: john at arbash-meinel.com-20090908210608-4fs0t0s43smepnzo
parent: john at arbash-meinel.com-20090908204304-bvrvpt10o8ux0e6z
committer: John Arbash Meinel <john at arbash-meinel.com>
branch nick: 2.1-memory-consumption
timestamp: Tue 2009-09-08 16:06:08 -0500
message:
  Add a get_key() function, which then returns tuples for the given items.
-------------- next part --------------
=== modified file 'bzrlib/_keys_type_c.c'
--- a/bzrlib/_keys_type_c.c	2009-09-08 20:43:04 +0000
+++ b/bzrlib/_keys_type_c.c	2009-09-08 21:06:08 +0000
@@ -32,12 +32,29 @@
     PyObject_HEAD
     unsigned char key_width;
     unsigned char num_keys; /* Not a Py_ssize_t like most containers */
-    PyStringObject *key_strings[1];
+    PyStringObject *key_strings[1]; /* key_width * num_keys entries */
 } Keys;
 
 /* Forward declaration */
 extern PyTypeObject KeysType;
 
+static void
+Keys_dealloc(Keys *keys)
+{
+    /* Do we want to use the Py_TRASHCAN_SAFE_BEGIN/END operations? */
+    if (keys->num_keys > 0) {
+        /* tuple deallocs from the end to the beginning. Not sure why, but
+         * we'll do the same here.
+         */
+        int i;
+        for(i = keys->num_keys - 1; i >= 0; --i) {
+            Py_XDECREF(keys->key_strings[i]);
+        }
+    }
+    Py_TYPE(keys)->tp_free((PyObject *)keys);
+}
+
+
 static PyObject *
 Keys_new(PyTypeObject *type, PyObject *args, PyObject *kwds)
 {
@@ -104,7 +121,7 @@
             return NULL;
         }
         Py_INCREF(obj);
-        self->key_strings[i] = obj;
+        self->key_strings[i] = (PyStringObject *)obj;
     }
     return (PyObject *)self;
 }
@@ -112,7 +129,42 @@
 static char Keys_doc[] =
     "C implementation of a Keys structure";
 
+static PyObject *
+Keys_get_key(Keys *self, PyObject *args) {
+    long offset;
+    long start;
+    long i;
+    PyObject *tpl = NULL, *obj = NULL;
+
+    if (!PyArg_ParseTuple(args, "l", &offset)) {
+        return NULL;
+    }
+    if (offset >= self->num_keys) {
+        PyErr_SetString(PyExc_IndexError, "offset out of range");
+        return NULL;
+    }
+    tpl = PyTuple_New(self->key_width);
+    if (!tpl) {
+        /* Malloc failure */
+        return NULL;
+    }
+    start = offset * self->key_width;
+    for (i = 0; i < self->key_width; ++i) {
+        obj = (PyObject *)self->key_strings[start + i];
+        Py_INCREF(obj);
+        PyTuple_SET_ITEM(tpl, i, obj);
+    }
+    return tpl;
+}
+
+static char Keys_get_key_doc[] = "get_keys(offset)";
+
+
 static PyMethodDef Keys_methods[] = {
+    {"get_key",
+     (PyCFunction)Keys_get_key,
+     METH_VARARGS,
+     Keys_get_key_doc},
     {NULL, NULL} /* sentinel */
 };
 
@@ -122,9 +174,7 @@
     "Keys",                                      /* tp_name */
     sizeof(Keys) - sizeof(PyStringObject *),           /* tp_basicsize */
     sizeof(PyObject *),                          /* tp_itemsize */
-    // We probably need a custom dealloc so that it decrefs the referenced
-    // PyString objects
-    0, // (destructor)PatienceSequenceMatcher_dealloc, /* tp_dealloc */
+    (destructor)Keys_dealloc,                    /* tp_dealloc */
     0,                                           /* tp_print */
     0,                                           /* tp_getattr */
     0,                                           /* tp_setattr */

=== modified file 'bzrlib/tests/test__keys_type.py'
--- a/bzrlib/tests/test__keys_type.py	2009-09-08 20:43:04 +0000
+++ b/bzrlib/tests/test__keys_type.py	2009-09-08 21:06:08 +0000
@@ -16,6 +16,8 @@
 
 """Tests for the Keys type."""
 
+import sys
+
 from bzrlib import (
     # _keys_py,
     errors,
@@ -60,7 +62,7 @@
 class TestKeysType(tests.TestCase):
 
     def test_create(self):
-        t = self.module.Keys(1, 'foo', 'bar')
+        k = self.module.Keys(1, 'foo', 'bar')
 
     def test_create_bad_args(self):
         self.assertRaises(TypeError, self.module.Keys)
@@ -74,3 +76,29 @@
         # too many args
         self.assertRaises(ValueError, self.module.Keys, 1, *lots_of_args)
         self.assertRaises(TypeError, self.module.Keys, 1, 'foo', 10)
+
+    def test_create_and_del_correct_refcount(self):
+        s = 'my custom' + ' foo bar'
+        n_ref = sys.getrefcount(s)
+        k = self.module.Keys(1, s)
+        self.assertEqual(n_ref + 1, sys.getrefcount(s))
+        del k
+        self.assertEqual(n_ref, sys.getrefcount(s))
+
+    def test_get_key(self):
+        f = 'fo' + 'o'
+        k = self.module.Keys(1, f, 'bar')
+        self.assertEqual(('foo',), k.get_key(0))
+        self.assertEqual(('bar',), k.get_key(1))
+        self.assertRaises(IndexError, k.get_key, 2)
+        n_refs = sys.getrefcount(f)
+        f_key = k.get_key(0)
+        self.assertEqual(n_refs + 1, sys.getrefcount(f))
+        del f_key
+        self.assertEqual(n_refs, sys.getrefcount(f))
+
+    def test_get_wide_key(self):
+        k = self.module.Keys(2, 'foo', 'bar', 'baz', 'bing')
+        self.assertEqual(('foo', 'bar'), k.get_key(0))
+        self.assertEqual(('baz', 'bing'), k.get_key(1))
+        self.assertRaises(IndexError, k.get_key, 2)