Implement Branch.upstream getter and setter.

Thanks to @cholin for the help!

src/branch.c

@@ -31,7 +31,10 @@
 #include "reference.h"
 #include "utils.h"
 
+
 extern PyObject *GitError;
+extern PyTypeObject ReferenceType;
+
 
 PyDoc_STRVAR(Branch_delete__doc__,
   "delete()\n"
@@ -161,6 +164,57 @@ PyObject* Branch_remote_name__get__(Branch *self)
 }
 
 
+PyDoc_STRVAR(Branch_upstream__doc__,
+  "The branch supporting the remote tracking branch or None if this is not a "
+  "remote tracking branch. Set to None to unset.");
+
+PyObject* Branch_upstream__get__(Branch *self)
+{
+    int err;
+    git_reference *c_reference;
+
+    CHECK_REFERENCE(self);
+
+    err = git_branch_upstream(&c_reference, self->reference);
+    if (err == GIT_ENOTFOUND)
+        Py_RETURN_NONE;
+    else if (err < GIT_OK)
+        return Error_set(err);
+
+    return wrap_branch(c_reference, self->repo);
+}
+
+int Branch_upstream__set__(Branch *self, Reference *py_ref)
+{
+    int err;
+    const char *branch_name = NULL;
+
+    CHECK_REFERENCE_INT(self);
+
+    if ((PyObject *)py_ref != Py_None) {
+        if (!PyObject_TypeCheck(py_ref, (PyTypeObject *)&ReferenceType)) {
+            PyErr_SetObject(PyExc_TypeError, (PyObject *)py_ref);
+            return -1;
+        }
+
+        CHECK_REFERENCE_INT(py_ref);
+        err = git_branch_name(&branch_name, py_ref->reference);
+        if (err < GIT_OK) {
+            Error_set(err);
+            return -1;
+        }
+    }
+
+    err = git_branch_set_upstream(self->reference, branch_name);
+    if (err < GIT_OK) {
+        Error_set(err);
+        return -1;
+    }
+
+    return 0;
+}
+
+
 PyMethodDef Branch_methods[] = {
     METHOD(Branch, delete, METH_NOARGS),
     METHOD(Branch, is_head, METH_NOARGS),
@@ -171,6 +225,7 @@ PyMethodDef Branch_methods[] = {
 PyGetSetDef Branch_getseters[] = {
     GETTER(Branch, branch_name),
     GETTER(Branch, remote_name),
+    GETSET(Branch, upstream),
     {NULL}
 };
 

test/test_branch.py

@@ -146,6 +146,24 @@ class BranchesEmptyRepoTestCase(utils.EmptyRepoTestCase):
                                          pygit2.GIT_BRANCH_REMOTE)
         self.assertEqual(branch.remote_name, 'origin')
 
+    def test_branch_upstream(self):
+        self.repo.remotes[0].fetch()
+        remote_master = self.repo.lookup_branch('origin/master',
+                                                pygit2.GIT_BRANCH_REMOTE)
+        master = self.repo.create_branch('master',
+                                         self.repo[remote_master.target.hex])
+
+        self.assertTrue(master.upstream is None)
+        master.upstream = remote_master
+        self.assertEqual(master.upstream.branch_name, 'origin/master')
+
+        def set_bad_upstream():
+            master.upstream = 2.5
+        self.assertRaises(TypeError, set_bad_upstream)
+
+        master.upstream = None
+        self.assertTrue(master.upstream is None)
+
 
 if __name__ == '__main__':
     unittest.main()