GH-105808: Fix a regression introduced in GH-101251 (#105910)

Yhg1s · web-flow · commit 1858db7cbdbf · 2023-06-19T17:09:04.000Z
Fix a regression introduced in GH-101251, causing GzipFile.flush() to not flush the compressor (nor pass along the zip_mode argument).
diff --git a/Lib/gzip.py b/Lib/gzip.py
@@ -370,8 +370,9 @@ def close(self):
     def flush(self,zlib_mode=zlib.Z_SYNC_FLUSH):
         self._check_not_closed()
         if self.mode == WRITE:
-            # Ensure the compressor's buffer is flushed
             self._buffer.flush()
+            # Ensure the compressor's buffer is flushed
+            self.fileobj.write(self.compress.flush(zlib_mode))
             self.fileobj.flush()
 
     def fileno(self):
diff --git a/Lib/test/test_gzip.py b/Lib/test/test_gzip.py
@@ -9,6 +9,7 @@
 import struct
 import sys
 import unittest
+import zlib
 from subprocess import PIPE, Popen
 from test.support import import_helper
 from test.support import os_helper
@@ -616,6 +617,54 @@ def test_issue44439(self):
             self.assertEqual(f.write(q), LENGTH)
             self.assertEqual(f.tell(), LENGTH)
 
+    def test_flush_flushes_compressor(self):
+        # See issue GH-105808.
+        b = io.BytesIO()
+        message = b"important message here."
+        with gzip.GzipFile(fileobj=b, mode='w') as f:
+            f.write(message)
+            f.flush()
+            partial_data = b.getvalue()
+        full_data = b.getvalue()
+        self.assertEqual(gzip.decompress(full_data), message)
+        # The partial data should contain the gzip header and the complete
+        # message, but not the end-of-stream markers (so we can't just
+        # decompress it directly).
+        with self.assertRaises(EOFError):
+            gzip.decompress(partial_data)
+        d = zlib.decompressobj(wbits=-zlib.MAX_WBITS)
+        f = io.BytesIO(partial_data)
+        gzip._read_gzip_header(f)
+        read_message = d.decompress(f.read())
+        self.assertEqual(read_message, message)
+
+    def test_flush_modes(self):
+        # Make sure the argument to flush is properly passed to the
+        # zlib.compressobj; see issue GH-105808.
+        class FakeCompressor:
+            def __init__(self):
+                self.modes = []
+            def compress(self, data):
+                return b''
+            def flush(self, mode=-1):
+                self.modes.append(mode)
+                return b''
+        b = io.BytesIO()
+        fc = FakeCompressor()
+        with gzip.GzipFile(fileobj=b, mode='w') as f:
+            f.compress = fc
+            f.flush()
+            f.flush(50)
+            f.flush(zlib_mode=100)
+        # The implicit close will also flush the compressor.
+        expected_modes = [
+            zlib.Z_SYNC_FLUSH,
+            50,
+            100,
+            -1,
+        ]
+        self.assertEqual(fc.modes, expected_modes)
+
 
 class TestOpen(BaseTest):
     def test_binary_modes(self):
diff --git a/Misc/NEWS.d/next/Library/2023-06-19-11-31-55.gh-issue-105808.NL-quu.rst b/Misc/NEWS.d/next/Library/2023-06-19-11-31-55.gh-issue-105808.NL-quu.rst
@@ -0,0 +1 @@
+Fix a regression introduced in GH-101251 for 3.12, causing :meth:`gzip.GzipFile.flush` to not flush the compressor (nor pass along the ``zip_mode`` argument).
