Do not create a zlib compressor when IGzipFile opens for writing

CHANGELOG.rst

@@ -9,6 +9,13 @@ Changelog
 
 version 1.10.0-dev
 ------------------
++ Speed up opening ``IGzipFile`` and ``igzip.open`` for writing. They no
+  longer create a 256 KiB zlib compressor only to replace it with an ISA-L
+  one. On Linux, glibc could return the compressor's memory to the system after
+  each file and page-fault it in again for the next one, so writing many small
+  gzip streams to ``io.BytesIO`` can be up to 5 times faster. On Python 3.11,
+  ``IGzipFile`` now closes a file that it opened when writing the gzip header
+  fails, as ``gzip.GzipFile`` does on newer Python versions.
 + Fix ``python -m isal.igzip`` replacing the input file with empty output and
   exiting with status 0 when the output file is the input file, as in
   ``-f -o data data``. It exits with an error now.

src/isal/igzip.py

@@ -156,7 +156,13 @@ def __init__(self, filename=None, mode=None,
                 f"{isal_zlib.ISAL_BEST_SPEED} and "
                 f"{isal_zlib.ISAL_BEST_COMPRESSION}, got {compresslevel}."
             )
+        # GzipFile.__init__ would create a zlib compressor only to discard it.
+        if (isinstance(mode, str) and mode.startswith(("w", "a", "x"))
+                and "t" not in mode and "U" not in mode):
+            self._init_writer(filename, mode, compresslevel, fileobj, mtime)
+            return
         super().__init__(filename, mode, compresslevel, fileobj, mtime)
+        # A write mode can also come from fileobj.
         if self.mode == WRITE:
             self.compress = isal_zlib.compressobj(compresslevel,
                                                   isal_zlib.DEFLATED,
@@ -167,6 +173,46 @@ def __init__(self, filename=None, mode=None,
             raw = _GzipReader(self.fileobj, READ_BUFFER_SIZE)
             self._buffer = io.BufferedReader(raw)
 
+    def _init_writer(self, filename, mode, compresslevel, fileobj, mtime):
+        # GzipFile.__init__'s write branch from Python 3.14, with isal_zlib.
+        # __del__ and close() need these even if opening fails.
+        self.mode = None
+        self.fileobj = None
+        self._buffer = None
+        if "b" not in mode:
+            mode += "b"
+        try:
+            if fileobj is None:
+                fileobj = self.myfileobj = builtins.open(filename, mode)
+            if filename is None:
+                filename = getattr(fileobj, "name", "")
+                if not isinstance(filename, (str, bytes)):
+                    filename = ""
+            else:
+                filename = os.fspath(filename)
+            self.mode = WRITE
+            self._init_write(filename)
+            self.compress = isal_zlib.compressobj(compresslevel,
+                                                  isal_zlib.DEFLATED,
+                                                  -isal_zlib.MAX_WBITS,
+                                                  isal_zlib.DEF_MEM_LEVEL,
+                                                  0)
+            self._write_mtime = mtime
+            if _GZIP_BUFFERS_WRITES:
+                self._buffer_size = gzip._WRITE_BUFFER_SIZE
+                self._buffer = io.BufferedWriter(
+                    gzip._WriteBufferStream(self),
+                    buffer_size=self._buffer_size)
+            self.fileobj = fileobj
+            self._write_gzip_header(compresslevel)
+        except BaseException:
+            self.fileobj = None
+            myfileobj = self.myfileobj
+            if myfileobj is not None:
+                self.myfileobj = None
+                myfileobj.close()
+            raise
+
     def readline(self, size=-1):
         """Read a line of decompressed data, or at most size bytes of it."""
         # GzipFile.closed tests whether fileobj is None, so testing fileobj

tests/test_igzip.py

@@ -9,6 +9,7 @@
  from CPython. Uses pytest which is easier to work with. Meant to complement
  the gzip module compliance tests. It should improve coverage as well."""
 
+import gc
 import gzip
 import io
 import itertools
@@ -60,6 +61,140 @@ def recording_compressobj(*args, **kwargs):
                        isal_zlib.DEF_MEM_LEVEL, 0), {})]
 
 
+@pytest.mark.parametrize("mode", ["w", "a", "x"])
+def test_igzipfile_writer_skips_zlib_compressor(monkeypatch, tmp_path, mode):
+    calls = []
+    monkeypatch.setattr(zlib, "compressobj",
+                        lambda *args, **kwargs: calls.append(args))
+    with igzip.IGzipFile(tmp_path / "test.gz", mode) as gzip_file:
+        gzip_file.write(DATA)
+    assert calls == []
+    assert gzip.decompress((tmp_path / "test.gz").read_bytes()) == DATA
+
+
+def test_igzipfile_writer_has_the_attributes_of_gzipfile():
+    # Catches attributes that new Python versions add to GzipFile.__init__.
+    with gzip.GzipFile(fileobj=io.BytesIO(), mode="wb") as gzip_file:
+        expected = {name: type(value)
+                    for name, value in vars(gzip_file).items()}
+    with igzip.IGzipFile(fileobj=io.BytesIO(), mode="wb") as igzip_file:
+        found = {name: type(value)
+                 for name, value in vars(igzip_file).items()}
+    assert found["compress"] is isal_zlib.Compress
+    del expected["compress"]
+    assert {name: found.get(name) for name in expected} == expected
+
+
+@pytest.mark.parametrize("mode", ["w", "wb", "a", "ab", "x", "xb"])
+@pytest.mark.parametrize("to_filename", [str, os.fsencode, Path])
+def test_igzipfile_write_to_filename(tmp_path, mode, to_filename):
+    filename = to_filename(tmp_path / "test.gz")
+    with igzip.IGzipFile(filename, mode, mtime=0) as gzip_file:
+        assert gzip_file.mode == igzip.WRITE
+        assert gzip_file.name == os.fspath(filename)
+        gzip_file.write(DATA)
+    compressed = (tmp_path / "test.gz").read_bytes()
+    assert compressed[3] == FNAME
+    assert compressed[10:15] == b"test\x00"
+    assert gzip.decompress(compressed) == DATA
+
+
+@pytest.mark.parametrize(["filename", "fileobj_name", "name", "header_name"], [
+    (None, "test.gz", "test.gz", b"test"),
+    (None, b"test.gz", b"test.gz", b"test"),
+    (None, 3, "", b""),
+    ("other.gz", "test.gz", "other.gz", b"other"),
+])
+def test_igzipfile_write_to_fileobj(filename, fileobj_name, name,
+                                    header_name):
+    fileobj = io.BytesIO()
+    fileobj.name = fileobj_name
+    with igzip.IGzipFile(filename, "wb", fileobj=fileobj,
+                         mtime=0) as gzip_file:
+        assert gzip_file.name == name
+        gzip_file.write(DATA)
+    assert not fileobj.closed
+    compressed = fileobj.getvalue()
+    if header_name:
+        assert compressed[3] == FNAME
+        assert compressed[10:].startswith(header_name + b"\x00")
+    else:
+        assert compressed[3] == 0
+    assert gzip.decompress(compressed) == DATA
+
+
+def test_igzipfile_append(tmp_path):
+    path = tmp_path / "test.gz"
+    for _ in range(2):
+        with igzip.IGzipFile(path, "ab") as gzip_file:
+            gzip_file.write(DATA)
+    assert gzip.decompress(path.read_bytes()) == DATA * 2
+
+
+def test_igzipfile_write_mode_from_fileobj(tmp_path):
+    with open(tmp_path / "test.gz", "wb") as fileobj:
+        with pytest.warns(FutureWarning):
+            gzip_file = igzip.IGzipFile(fileobj=fileobj)
+        with gzip_file:
+            assert gzip_file.mode == igzip.WRITE
+            assert isinstance(gzip_file.compress, isal_zlib.Compress)
+            gzip_file.write(DATA)
+    assert gzip.decompress((tmp_path / "test.gz").read_bytes()) == DATA
+
+
+@pytest.mark.parametrize(["filename", "mode", "error"], [
+    (None, "wb", TypeError),
+    ("missing/test.gz", "wb", FileNotFoundError),
+    ("missing/test.gz", "ab", FileNotFoundError),
+    ("exists.gz", "xb", FileExistsError),
+])
+def test_igzipfile_write_open_error(monkeypatch, tmp_path, filename, mode,
+                                    error):
+    (tmp_path / "exists.gz").write_bytes(b"data")
+    if filename is not None:
+        filename = tmp_path / filename
+    # Keep earlier tests' unraisable errors out of the hook.
+    gc.collect()
+    unraisable = []
+    monkeypatch.setattr(sys, "unraisablehook", unraisable.append)
+    with pytest.raises(error):
+        igzip.IGzipFile(filename, mode)
+    gc.collect()
+    assert unraisable == []
+    assert (tmp_path / "exists.gz").read_bytes() == b"data"
+
+
+def test_igzipfile_mode_not_str(monkeypatch):
+    gc.collect()
+    unraisable = []
+    monkeypatch.setattr(sys, "unraisablehook", unraisable.append)
+    with pytest.raises(TypeError):
+        igzip.IGzipFile(fileobj=io.BytesIO(), mode=b"wb")
+    gc.collect()
+    assert unraisable == []
+
+
+@pytest.mark.parametrize("mode", ["wb", "ab", "xb"])
+def test_igzipfile_header_error_closes_opened_file(monkeypatch, tmp_path,
+                                                   mode):
+    fileobjs = []
+
+    def failing_write_gzip_header(self, compresslevel):
+        fileobjs.append(self.fileobj)
+        raise OSError("No space left on device")
+
+    monkeypatch.setattr(igzip.IGzipFile, "_write_gzip_header",
+                        failing_write_gzip_header)
+    with pytest.raises(OSError, match="No space left on device"):
+        igzip.IGzipFile(tmp_path / "test.gz", mode)
+    fileobj = io.BytesIO()
+    with pytest.raises(OSError, match="No space left on device"):
+        igzip.IGzipFile(fileobj=fileobj, mode=mode)
+    assert fileobjs[0].closed
+    assert fileobjs[1] is fileobj
+    assert not fileobj.closed
+
+
 def test_write_readonly_file():
     with igzip.IGzipFile(TEST_FILE, "rb") as test:
         with pytest.raises(OSError) as error: