From f5082b320b595bf4264e9a26d3dccded9fb83803 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?J=2E=20David=20Ib=C3=A1=C3=B1ez?= <jdavid.ibp@gmail.com>
Date: Fri, 3 May 2013 23:31:43 +0200
Subject: [PATCH] Add "len(treebuider)" to the API

---
 src/treebuilder.c        | 28 +++++++++++++++++++---------
 test/test_treebuilder.py | 10 +++++++++-
 2 files changed, 28 insertions(+), 10 deletions(-)

diff --git a/src/treebuilder.c b/src/treebuilder.c
index 30712aa..4329a4d 100644
--- a/src/treebuilder.c
+++ b/src/treebuilder.c
@@ -56,20 +56,16 @@ TreeBuilder_insert(TreeBuilder *self, PyObject *args)
     git_oid oid;
     const char *fname;
 
-    if (!PyArg_ParseTuple(args, "sOi", &fname, &py_oid, &attr)) {
+    if (!PyArg_ParseTuple(args, "sOi", &fname, &py_oid, &attr))
         return NULL;
-    }
 
     len = py_str_to_git_oid(py_oid, &oid);
-    if (len < 0) {
+    if (len < 0)
         return NULL;
-    }
 
     err = git_treebuilder_insert(NULL, self->bld, fname, &oid, attr);
-    if (err < 0) {
-        Error_set(err);
-        return NULL;
-    }
+    if (err < 0)
+        return Error_set(err);
 
     Py_RETURN_NONE;
 }
@@ -138,6 +134,20 @@ PyMethodDef TreeBuilder_methods[] = {
 };
 
 
+Py_ssize_t
+TreeBuilder_len(TreeBuilder *self)
+{
+    return (Py_ssize_t)git_treebuilder_entrycount(self->bld);
+}
+
+
+PyMappingMethods TreeBuilder_as_mapping = {
+    (lenfunc)TreeBuilder_len,     /* mp_length */
+    0,                            /* mp_subscript */
+    0,                            /* mp_ass_subscript */
+};
+
+
 PyDoc_STRVAR(TreeBuilder__doc__, "TreeBuilder objects.");
 
 PyTypeObject TreeBuilderType = {
@@ -153,7 +163,7 @@ PyTypeObject TreeBuilderType = {
     0,                                         /* tp_repr           */
     0,                                         /* tp_as_number      */
     0,                                         /* tp_as_sequence    */
-    0,                                         /* tp_as_mapping     */
+    &TreeBuilder_as_mapping,                   /* tp_as_mapping     */
     0,                                         /* tp_hash           */
     0,                                         /* tp_call           */
     0,                                         /* tp_str            */
diff --git a/test/test_treebuilder.py b/test/test_treebuilder.py
index 161d400..2c30f76 100644
--- a/test/test_treebuilder.py
+++ b/test/test_treebuilder.py
@@ -41,25 +41,33 @@ class TreeBuilderTest(utils.BareRepoTestCase):
     def test_new_empty_treebuilder(self):
         self.repo.TreeBuilder()
 
+
     def test_noop_treebuilder(self):
         tree = self.repo[TREE_SHA]
         bld = self.repo.TreeBuilder(TREE_SHA)
         result = bld.write()
+
+        self.assertEqual(len(bld), len(tree))
         self.assertEqual(tree.oid, result)
 
+
     def test_noop_treebuilder_from_tree(self):
         tree = self.repo[TREE_SHA]
         bld = self.repo.TreeBuilder(tree)
         result = bld.write()
+
+        self.assertEqual(len(bld), len(tree))
         self.assertEqual(tree.oid, result)
 
+
     def test_rebuild_treebuilder(self):
         tree = self.repo[TREE_SHA]
         bld = self.repo.TreeBuilder()
         for e in tree:
             bld.insert(e.name, e.hex, e.filemode)
-
         result = bld.write()
+
+        self.assertEqual(len(bld), len(tree))
         self.assertEqual(tree.oid, result)