Faster gzip.compress implementation

python · ambv · Sep 2, 2021 · Mar 24, 2021 · Mar 24, 2021 · Mar 24, 2021
commit d1c86dcf1dde1fb8442c93d999d2058c0cc7a4b1
diff --git a/Lib/gzip.py b/Lib/gzip.py
@@ -544,14 +544,43 @@ def _rewind(self):
         super()._rewind()
         self._new_member = True
 
-def compress(data, compresslevel=_COMPRESS_LEVEL_BEST, *, mtime=None):
+
+def _create_simple_gzip_header(compresslevel: int,
+                               mtime = None) -> bytes:
+    """
+    Write a simple gzip header with no extra fields.
+    :param compresslevel: Compresslevel used to determine the xfl bytes.
+    :param mtime: The mtime (must support conversion to a 32-bit integer).
+    :return: A bytes object representing the gzip header.
+    """
+    if mtime is None:
+        mtime = time.time()
+    if compresslevel == _COMPRESS_LEVEL_BEST:
+        xfl = 2
+    elif compresslevel == _COMPRESS_LEVEL_FAST:
+        xfl = 4
+    else:
+        xfl = 0
+    # Pack ID1 and ID2 magic bytes, method (8=deflate), header flags (no extra
+    # fields added to header), mtime, xfl and os (255 for unknown OS).
+    return struct.pack("<BBBBLBB", 0x1f, 0x8b, 8, 0, int(mtime), xfl, 255)
+
+
+def compress(data, compresslevel=_COMPRESS_LEVEL_BEST, *, mtime=0):
     """Compress data in one shot and return the compressed string.
     Optional argument is the compression level, in range of 0-9.
+    mtime can be used to set the modification time. 0 is default,
+    use 'None' for current time.
     """
-    buf = io.BytesIO()
-    with GzipFile(fileobj=buf, mode='wb', compresslevel=compresslevel, mtime=mtime) as f:
-        f.write(data)
-    return buf.getvalue()
+    if mtime == 0:
+        # Use zlib as it creates the header with 0 mtime by default.
+        # This is faster and with less overhead.
+        return zlib.compress(data, level=compresslevel, wbits=31)
+    header = _create_simple_gzip_header(compresslevel, mtime)
+    trailer = struct.pack("LL", zlib.crc32(data), (len(data) & 0xffffffff))
+    # Wbits=-15 creates a raw deflate block.
+    return b"".join([header, zlib.compress(data, wbits=-15), trailer])
+
 
 def decompress(data):
     """Decompress a gzip compressed string in one shot.