Make TreeBuilder grow out of the repository

A TreeBuilder is a lot more useful (and easy to use) when it belongs
to a repository.

pygit2.c

@@ -937,6 +937,45 @@ Repository_status_file(Repository *self, PyObject *value)
     return PyInt_FromLong(status);
 }
 
+static PyObject *
+Repository_TreeBuilder(Repository *self, PyObject *args)
+{
+    TreeBuilder *builder;
+    git_treebuilder *bld;
+    PyObject *py_oid = NULL;
+    size_t oid_len;
+    git_oid oid;
+    git_tree *tree = NULL;
+    int err;
+
+    if (!PyArg_ParseTuple(args, "|O", &py_oid))
+        return NULL;
+
+    if (py_oid) {
+        oid_len = py_str_to_git_oid(py_oid, &oid);
+        TODO_SUPPORT_SHORT_HEXS(oid_len)
+            if (oid_len == 0)
+                return NULL;
+
+        err = git_tree_lookup(&tree, self->repo, &oid);
+        if (err < 0)
+            return Error_set(err);
+    }
+
+    err = git_treebuilder_create(&bld, tree);
+    if (err < 0)
+        return Error_set(err);
+
+    builder = PyObject_New(TreeBuilder, &TreeBuilderType);
+    if (builder) {
+        builder->repo = self;
+        builder->bld = bld;
+        Py_INCREF(self);
+    }
+
+    return (PyObject*)builder;
+}
+
 static PyMethodDef Repository_methods[] = {
     {"create_commit", (PyCFunction)Repository_create_commit, METH_VARARGS,
      "Create a new commit object, return its SHA."},
@@ -970,6 +1009,8 @@ static PyMethodDef Repository_methods[] = {
      "as keys and status flags as values.\nSee pygit2.GIT_STATUS_*."},
     {"status_file", (PyCFunction)Repository_status_file, METH_O,
      "Returns the status of the given file path."},
+    {"TreeBuilder", (PyCFunction)Repository_TreeBuilder, METH_VARARGS,
+     "Create a TreeBuilder object for this repository."},
     {NULL}
 };
 
@@ -1578,36 +1619,10 @@ static PyTypeObject TreeType = {
     0,                                         /* tp_new            */
 };
 
-static int
-TreeBuilder_init(TreeBuilder *self, PyObject *args, PyObject *kwds)
-{
-    Tree *py_tree = NULL;
-    git_tree *tree;
-    int err;
-
-    if (kwds) {
-        PyErr_SetString(PyExc_TypeError,
-                        "TreeBuilder takes no keyword arguments");
-        return -1;
-    }
-
-    if (!PyArg_ParseTuple(args, "|O", &py_tree))
-        return -1;
-
-    tree = py_tree == NULL ? NULL : py_tree->tree;
-
-    err = git_treebuilder_create(&self->bld, tree);
-    if (err < 0) {
-        Error_set(err);
-        return -1;
-    }
-
-    return 0;
-}
-
 static void
 TreeBuilder_dealloc(TreeBuilder* self)
 {
+    Py_XDECREF(self->repo);
     git_treebuilder_free(self->bld);
     PyObject_Del(self);
 }
@@ -1633,13 +1648,12 @@ TreeBuilder_insert(TreeBuilder *self, TreeEntry *py_tentry)
 }
 
 static PyObject *
-TreeBuilder_write(TreeBuilder *self, Repository *py_repo)
+TreeBuilder_write(TreeBuilder *self)
 {
     int err;
     git_oid oid;
-    git_repository *repo = py_repo->repo;
 
-    err = git_treebuilder_write(&oid, repo, self->bld);
+    err = git_treebuilder_write(&oid, self->repo->repo, self->bld);
     if (err < 0)
         return Error_set(err);
 
@@ -1649,7 +1663,7 @@ TreeBuilder_write(TreeBuilder *self, Repository *py_repo)
 static PyMethodDef TreeBuilder_methods[] = {
     {"insert", (PyCFunction)TreeBuilder_insert, METH_O,
      "Insert or replace an entry in the treebuilder"},
-    {"write", (PyCFunction)TreeBuilder_write, METH_O,
+    {"write", (PyCFunction)TreeBuilder_write, METH_NOARGS,
      "Write the tree to the given repository"},
     {NULL, NULL, 0, NULL}
 };
@@ -1690,7 +1704,7 @@ static PyTypeObject TreeBuilderType = {
     0,                                         /* tp_descr_get      */
     0,                                         /* tp_descr_set      */
     0,                                         /* tp_dictoffset     */
-    (initproc)TreeBuilder_init,                /* tp_init           */
+    0,                                         /* tp_init           */
     0,                                         /* tp_alloc          */
     0,                                         /* tp_new            */
 };
@@ -3073,9 +3087,6 @@ moduleinit(PyObject* m)
     Py_INCREF(&TreeType);
     PyModule_AddObject(m, "Tree", (PyObject *)&TreeType);
 
-    Py_INCREF(&TreeBuilderType);
-    PyModule_AddObject(m, "TreeBuilder", (PyObject *)&TreeBuilderType);
-
     Py_INCREF(&BlobType);
     PyModule_AddObject(m, "Blob", (PyObject *)&BlobType);
 

test/test_treebuilder.py

@@ -42,21 +42,21 @@ TREE_SHA = '967fce8df97cc71722d3c2a5930ef3e6f1d27b12'
 
 class TreeBuilderTest(utils.BareRepoTestCase):
     def test_new_empty_treebuilder(self):
-        bld = pygit2.TreeBuilder()
+        bld = self.repo.TreeBuilder()
 
     def test_noop_treebuilder(self):
         tree = self.repo[TREE_SHA]
-        bld = pygit2.TreeBuilder(tree)
-        result = bld.write(self.repo)
+        bld = self.repo.TreeBuilder(TREE_SHA)
+        result = bld.write()
         self.assertEqual(tree.oid, result)
 
     def test_rebuild_treebuilder(self):
         tree = self.repo[TREE_SHA]
-        bld = pygit2.TreeBuilder(tree)
+        bld = self.repo.TreeBuilder(TREE_SHA)
         for e in tree:
             bld.insert(e)
 
-        result = bld.write(self.repo)
+        result = bld.write()
         self.assertEqual(tree.oid, result)
 
 if __name__ == '__main__':