object: allow passing a type object to peel()

Allow usage of a python type instead of having to use the libgit2 constants.
folpindo · Feb 11, 2014 · b2cd25c · b2cd25c
1 parent de3dba6
commit b2cd25c
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 5 deletions.
diff --git a/src/object.c b/src/object.c
@@ -127,6 +127,24 @@ Object_read_raw(Object *self)
  return aux;
 }
 
+static git_otype
+py_type_to_git_type(PyTypeObject *py_type)
+{
+ git_otype type = GIT_OBJ_BAD;
+
+ if (py_type == &CommitType) {
+ type = GIT_OBJ_COMMIT;
+ } else if (py_type == &TreeType) {
+ type = GIT_OBJ_TREE;
+ } else if (py_type == &BlobType) {
+ type = GIT_OBJ_BLOB;
+ } else if (py_type == &TagType) {
+ type = GIT_OBJ_TAG;
+ }
+
+ return type;
+}
+
 PyDoc_STRVAR(Object_peel__doc__,
  "peel(target_type) -> Object\n"
  "\n"
@@ -135,12 +153,21 @@ PyDoc_STRVAR(Object_peel__doc__,
 PyObject *
 Object_peel(Object *self, PyObject *py_type)
 {
- int type, err;
+ int type = -1, err;
  git_object *peeled;
 
- type = PyLong_AsLong(py_type);
- if (type == -1 && PyErr_Occurred())
+ if (PyLong_Check(py_type)) {
+ type = PyLong_AsLong(py_type);
+ if (type == -1 && PyErr_Occurred())
+ return NULL;
+ } else if (PyType_Check(py_type)) {
+ type = py_type_to_git_type((PyTypeObject *) py_type);
+ }
+
+ if (type == -1) {
+ PyErr_SetString(PyExc_ValueError, "invalid target type");
  return NULL;
+ }
 
  err = git_object_peel(&peeled, self->obj, (git_otype)type);
  if (err < 0)

diff --git a/test/test_object.py b/test/test_object.py
@@ -33,7 +33,7 @@
 import unittest
 
 import pygit2
-from pygit2 import GIT_OBJ_TREE, GIT_OBJ_TAG
+from pygit2 import GIT_OBJ_TREE, GIT_OBJ_TAG, Tree, Tag
 from . import utils
 
 
@@ -54,14 +54,29 @@ def test_peel_commit(self):
  # and peel to the tree
  tree = commit.peel(GIT_OBJ_TREE)
 
- self.assertEqual(type(tree), pygit2.Tree)
+ self.assertEqual(type(tree), Tree)
  self.assertEqual(str(tree.id), 'fd937514cb799514d4b81bb24c5fcfeb6472b245')
 
+ def test_peel_commit_type(self):
+ commit_id = self.repo.lookup_reference('refs/heads/master').target
+ commit = self.repo[commit_id]
+ tree = commit.peel(Tree)
+
+ self.assertEqual(type(tree), Tree)
+ self.assertEqual(str(tree.id), 'fd937514cb799514d4b81bb24c5fcfeb6472b245')
+
+
  def test_invalid(self):
  commit_id = self.repo.lookup_reference('refs/heads/master').target
  commit = self.repo[commit_id]
 
  self.assertRaises(ValueError, commit.peel, GIT_OBJ_TAG)
 
+ def test_invalid_type(self):
+ commit_id = self.repo.lookup_reference('refs/heads/master').target
+ commit = self.repo[commit_id]
+
+ self.assertRaises(ValueError, commit.peel, Tag)
+
 if __name__ == '__main__':
  unittest.main()