TreeEntry: add .blob and .tree properties

rcoup · rcoup · commit 8215bc9212f4 · 2019-06-14T15:56:45.000+01:00
Saves doing another `repo[tree_entry.id]` lookup

Tree navigation: `blob = (tree / 'path' / 'file').blob`
diff --git a/src/tree.c b/src/tree.c
@@ -31,6 +31,7 @@
 #include "error.h"
 #include "utils.h"
 #include "repository.h"
+#include "object.h"
 #include "oid.h"
 #include "tree.h"
 #include "diff.h"
@@ -177,25 +178,14 @@ TreeEntry_hex__get__(TreeEntry *self)
     return git_oid_to_py_str(git_tree_entry_id(self->entry));
 }
 
-PyObject *
-TreeEntry_repr(TreeEntry *self)
-{
-    char str[GIT_OID_HEXSZ + 1] = { 0 };
-    const char *typename;
-
-    typename = git_object_type2string(git_tree_entry_type(self->entry));
-    git_oid_fmt(str, git_tree_entry_id(self->entry));
-    return PyString_FromFormat("pygit2.TreeEntry('%s', %s, %s)", git_tree_entry_name(self->entry), typename, str);
-}
-
 git_tree *
 treeentry_to_subtree(TreeEntry* self)
 {
     Repository *py_repo;
     git_tree *subtree = NULL;
 
     if (git_tree_entry_type(self->entry) != GIT_OBJ_TREE) {
-        PyErr_SetString(PyExc_TypeError, "Only supported for trees");
+        PyErr_SetString(PyExc_TypeError, "Only for trees");
         return NULL;
     }
 
@@ -214,6 +204,62 @@ treeentry_to_subtree(TreeEntry* self)
     return subtree;
 }
 
+PyObject *
+treeentry_to_object(TreeEntry* self)
+{
+    Repository *py_repo;
+    git_object *obj= NULL;
+
+    if (self->repo == NULL) {
+        PyErr_SetString(PyExc_ValueError, "No repository associated with this TreeEntry");
+        return NULL;
+    }
+    py_repo = self->repo;
+
+    int err = git_tree_entry_to_object(&obj, py_repo->repo, self->entry);
+    if (err < 0) {
+        Error_set(err);
+        return NULL;
+    }
+
+    return wrap_object(obj, py_repo);
+}
+PyDoc_STRVAR(TreeEntry_tree__doc__, "Subtree");
+
+PyObject *
+TreeEntry_tree__get__(TreeEntry *self)
+{
+    git_tree* subtree = treeentry_to_subtree(self);
+    if (subtree == NULL)
+        return NULL;
+
+    return wrap_object((git_object*)subtree, self->repo);
+}
+
+PyDoc_STRVAR(TreeEntry_blob__doc__, "Blob");
+
+PyObject *
+TreeEntry_blob__get__(TreeEntry *self)
+{
+    if (git_tree_entry_type(self->entry) != GIT_OBJ_BLOB) {
+        PyErr_SetString(PyExc_TypeError, "Only for blobs");
+        return NULL;
+    }
+
+    return treeentry_to_object(self);
+}
+
+PyObject *
+TreeEntry_repr(TreeEntry *self)
+{
+    char str[GIT_OID_HEXSZ + 1] = { 0 };
+    const char *typename;
+
+    typename = git_object_type2string(git_tree_entry_type(self->entry));
+    git_oid_fmt(str, git_tree_entry_id(self->entry));
+    return PyString_FromFormat("pygit2.TreeEntry('%s', %s, %s)", git_tree_entry_name(self->entry), typename, str);
+}
+
 int
 TreeEntry_contains(TreeEntry *self, PyObject *py_name)
 {
@@ -295,6 +341,8 @@ PyGetSetDef TreeEntry_getseters[] = {
     GETTER(TreeEntry, id),
     GETTER(TreeEntry, hex),
     GETTER(TreeEntry, type),
+    GETTER(TreeEntry, blob),
+    GETTER(TreeEntry, tree),
     {NULL}
 };
 
diff --git a/src/tree.h b/src/tree.h
@@ -38,6 +38,8 @@ PyObject* TreeEntry_get_filemode(TreeEntry *self);
 PyObject* TreeEntry_get_name(TreeEntry *self);
 PyObject* TreeEntry_get_oid(TreeEntry *self);
 PyObject* TreeEntry_get_hex(TreeEntry *self);
+PyObject* TreeEntry_get_blob(TreeEntry *self);
+PyObject* TreeEntry_get_tree(TreeEntry *self);
 
 TreeEntry* tree_getitem_by_index(const git_tree *tree, Repository *repo, PyObject *py_index);
 TreeEntry* tree_getitem_by_path(const git_tree *tree, Repository *repo, PyObject *py_path);
diff --git a/test/test_tree.py b/test/test_tree.py
@@ -83,7 +83,10 @@ def test_read_tree(self):
         self.assertRaisesWithArg(KeyError, 'ab/cd', lambda: tree['ab/cd'])
         self.assertRaisesWithArg(KeyError, 'ab/cd', lambda: tree / 'ab/cd')
         self.assertRaisesWithArg(KeyError, 'ab', lambda: tree / 'c' / 'ab')
-        self.assertRaisesWithArg(TypeError, 'Only supported for trees', lambda: tree / 'a' / 'cd')
+        self.assertRaisesWithArg(TypeError, 'Only for trees', lambda: tree / 'a' / 'cd')
+
+        self.assertRaisesWithArg(TypeError, 'Only for trees', lambda: (tree / 'c' / 'd').tree)
+        self.assertRaisesWithArg(TypeError, 'Only for blobs', lambda: (tree / 'c').blob)
 
     def test_equality(self):
         tree_a = self.repo['18e2d2e9db075f9eb43bcb2daa65a2867d29a15e']
@@ -113,6 +116,8 @@ def test_read_subtree(self):
         sha = '297efb891a47de80be0cfe9c639e4b8c9b450989'
         self.assertTreeEntryEqual(subtree[0], sha, 'd', 0o0100644)
 
+        subtree_entry = (tree / 'c')
+        assert subtree_entry.tree == self.repo[subtree_entry.id]
 
     def test_new_tree(self):
         repo = self.repo
@@ -154,6 +159,10 @@ def test_new_tree(self):
         assert y.type == 'blob'
         assert z.type == 'tree'
 
+        assert x.blob == repo[x.id]
+        assert y.blob == repo[y.id]
+        assert z.tree == repo[z.id]
+
 
     def test_modify_tree(self):
         tree = self.repo[TREE_SHA]
