Add compression_level arguments in pack writing functions.

dulwich/pack.py

@@ -1531,13 +1531,14 @@ def pack_object_header(type_num, delta_base, size):
     return bytearray(header)
 
 
-def write_pack_object(f, type, object, sha=None):
+def write_pack_object(f, type, object, sha=None, compression_level=-1):
     """Write pack object to a file.
 
     Args:
       f: File to write to
       type: Numeric type of the object
       object: Object to write
+      compression_level: the zlib compression level
     Returns: Tuple with offset at which the object was written, and crc32
     """
     if type in DELTA_TYPES:
@@ -1545,7 +1546,7 @@ def write_pack_object(f, type, object, sha=None):
     else:
         delta_base = None
     header = bytes(pack_object_header(type, delta_base, len(object)))
-    comp_data = zlib.compress(object)
+    comp_data = zlib.compress(object, level=compression_level)
     crc32 = 0
     for data in (header, comp_data):
         f.write(data)
@@ -1555,7 +1556,8 @@ def write_pack_object(f, type, object, sha=None):
     return crc32 & 0xffffffff
 
 
-def write_pack(filename, objects, deltify=None, delta_window_size=None):
+def write_pack(filename, objects, deltify=None, delta_window_size=None,
+               compression_level=-1):
     """Write a new pack data file.
 
     Args:
@@ -1564,11 +1566,13 @@ def write_pack(filename, objects, deltify=None, delta_window_size=None):
         Should provide __len__
       window_size: Delta window size
       deltify: Whether to deltify pack objects
+      compression_level: the zlib compression level
     Returns: Tuple with checksum of pack file and index file
     """
     with GitFile(filename + '.pack', 'wb') as f:
         entries, data_sum = write_pack_objects(
-            f, objects, delta_window_size=delta_window_size, deltify=deltify)
+            f, objects, delta_window_size=delta_window_size, deltify=deltify,
+            compression_level=compression_level)
     entries = sorted([(k, v[0], v[1]) for (k, v) in entries.items()])
     with GitFile(filename + '.idx', 'wb') as f:
         return data_sum, write_pack_index_v2(f, entries, data_sum)
@@ -1632,7 +1636,8 @@ def pack_objects_to_data(objects):
              for (o, path) in objects))
 
 
-def write_pack_objects(f, objects, delta_window_size=None, deltify=None):
+def write_pack_objects(f, objects, delta_window_size=None, deltify=None,
+                       compression_level=-1):
     """Write a new pack data file.
 
     Args:
@@ -1642,6 +1647,7 @@ def write_pack_objects(f, objects, delta_window_size=None, deltify=None):
       window_size: Sliding window size for searching for deltas;
                         Set to None for default window size.
       deltify: Whether to deltify objects
+      compression_level: the zlib compression level to use
     Returns: Dict mapping id -> (offset, crc32 checksum), pack checksum
     """
     if deltify is None:
@@ -1654,10 +1660,13 @@ def write_pack_objects(f, objects, delta_window_size=None, deltify=None):
     else:
         pack_contents_count, pack_contents = pack_objects_to_data(objects)
 
-    return write_pack_data(f, pack_contents_count, pack_contents)
+    return write_pack_data(
+        f, pack_contents_count, pack_contents,
+        compression_level=compression_level)
 
 
-def write_pack_data(f, num_records, records, progress=None):
+def write_pack_data(
+        f, num_records, records, progress=None, compression_level=-1):
     """Write a new pack data file.
 
     Args:
@@ -1665,6 +1674,7 @@ def write_pack_data(f, num_records, records, progress=None):
       num_records: Number of records
       records: Iterator over type_num, object_id, delta_base, raw
       progress: Function to report progress to
+      compression_level: the zlib compression level
     Returns: Dict mapping id -> (offset, crc32 checksum), pack checksum
     """
     # Write the pack
@@ -1686,7 +1696,8 @@ def write_pack_data(f, num_records, records, progress=None):
             else:
                 type_num = OFS_DELTA
                 raw = (offset - base_offset, raw)
-        crc32 = write_pack_object(f, type_num, raw)
+        crc32 = write_pack_object(
+            f, type_num, raw, compression_level=compression_level)
         entries[object_id] = (offset, crc32)
     return entries, f.write_sha()
 

dulwich/tests/test_pack.py

@@ -591,6 +591,18 @@ class WritePackTests(TestCase):
         sha_b.update(f.getvalue()[offset:])
         self.assertEqual(sha_a.digest(), sha_b.digest())
 
+    def test_write_pack_object_compression_level(self):
+        f = BytesIO()
+        f.write(b'header')
+        offset = f.tell()
+        sha_a = sha1(b'foo')
+        sha_b = sha_a.copy()
+        write_pack_object(
+            f, Blob.type_num, b'blob', sha=sha_a, compression_level=6)
+        self.assertNotEqual(sha_a.digest(), sha_b.digest())
+        sha_b.update(f.getvalue()[offset:])
+        self.assertEqual(sha_a.digest(), sha_b.digest())
+
 
 pack_checksum = hex_to_sha('721980e866af9a5f93ad674144e1459b8ba3e7b7')