Add Repo.update_shallow.

dulwich/repo.py

@@ -468,6 +468,14 @@ class BaseRepo(object):
         """
         return set()
 
+    def update_shallow(self, new_shallow, new_unshallow):
+        """Update the list of shallow objects.
+
+        :param new_shallow: Newly shallow objects
+        :param new_unshallow: Newly no longer shallow objects
+        """
+        raise NotImplementedError(self.update_shallow)
+
     def get_peeled(self, ref):
         """Get the peeled value of a ref.
 
@@ -935,6 +943,16 @@ class Repo(BaseRepo):
         with f:
             return set(l.strip() for l in f)
 
+    def update_shallow(self, new_shallow, new_unshallow):
+        shallow = self.get_shallow()
+        if new_shallow:
+            shallow.update(new_shallow)
+        if new_unshallow:
+            shallow.difference_update(new_unshallow)
+        self._put_named_file(
+            'shallow',
+            b''.join([sha + b'\n' for sha in shallow]))
+
     def index_path(self):
         """Return path to the index file."""
         return os.path.join(self.controldir(), INDEX_FILENAME)

dulwich/tests/test_repository.py

@@ -662,13 +662,25 @@ class BuildRepoRootTests(TestCase):
         self.assertEqual([], r[commit_sha].parents)
         self._root_commit = commit_sha
 
-    def test_shallow(self):
+    def test_get_shallow(self):
         self.assertEqual(set(), self._repo.get_shallow())
         with open(os.path.join(self._repo.path, '.git', 'shallow'), 'wb') as f:
             f.write(b'a90fa2d900a17e99b433217e988c4eb4a2e9a097\n')
         self.assertEqual({b'a90fa2d900a17e99b433217e988c4eb4a2e9a097'},
                          self._repo.get_shallow())
 
+    def test_update_shallow(self):
+        self._repo.update_shallow(None, None)  # no op
+        self.assertEquals(set(), self._repo.get_shallow())
+        self._repo.update_shallow([b'a90fa2d900a17e99b433217e988c4eb4a2e9a097'],
+                                  None)
+        self.assertEqual({b'a90fa2d900a17e99b433217e988c4eb4a2e9a097'},
+                         self._repo.get_shallow())
+        self._repo.update_shallow([b'a90fa2d900a17e99b433217e988c4eb4a2e9a097'],
+                                  [b'f9e39b120c68182a4ba35349f832d0e4e61f485c'])
+        self.assertEqual({b'a90fa2d900a17e99b433217e988c4eb4a2e9a097'},
+                         self._repo.get_shallow())
+
     def test_build_repo(self):
         r = self._repo
         self.assertEqual(b'ref: refs/heads/master', r.refs.read_ref(b'HEAD'))