Merge pull request #61 from pycompression/release_0.5.1

Release 0.5.1

Author: rhpvorderman

CHANGELOG.rst

@@ -7,6 +7,13 @@ Changelog
 .. This document is user facing. Please word the changes in such a way
 .. that users understand how the changes affect the new version.
 
+version 0.5.1
+-----------------
++ Fix a bug where flushing in threaded mode did not write the data to the
+  output file.
++ Threaded reading and writing do no longer block exiting when an exception
+  occurs in the main thread.
+
 version 0.5.0
 -----------------
 + Wheels are now build for MacOS arm64 architectures.

docs/conf.py

@@ -4,26 +4,19 @@
 # list see the documentation:
 # https://www.sphinx-doc.org/en/master/usage/configuration.html
 
-from distutils.dist import DistributionMetadata
-from pathlib import Path
-
-import pkg_resources
+import zlib_ng
 
 # -- Project information -----------------------------------------------------
 
-# Get package information from the installed package.
-package = pkg_resources.get_distribution("zlib-ng")
-metadata_file = Path(package.egg_info) / Path(package.PKG_INFO)
-metadata = DistributionMetadata(path=str(metadata_file))
-
 project = 'python-zlib-ng'
 copyright = '2023, Leiden University Medical Center'
 author = 'Leiden University Medical Center'
 
+
 # The short X.Y version
-version = package.parsed_version.base_version
+version = zlib_ng.__version__
 # The full version, including alpha/beta/rc tags
-release = package.version
+release = zlib_ng.__version__
 
 
 # -- General configuration ---------------------------------------------------

requirements-docs.txt

@@ -1,4 +1,4 @@
 sphinx
 setuptools
-sphinx-rtd-theme>=1.2.0rc3,<1.3
+sphinx-rtd-theme
 sphinx-argparse

setup.py

@@ -151,6 +151,7 @@ def build_zlib_ng():
         "Programming Language :: Python :: 3.10",
         "Programming Language :: Python :: 3.11",
         "Programming Language :: Python :: 3.12",
+        "Programming Language :: Python :: 3.13",
         "Programming Language :: Python :: Implementation :: CPython",
         "Programming Language :: Python :: Implementation :: PyPy",
         "Programming Language :: C",

src/zlib_ng/gzip_ng_threaded.py

@@ -60,7 +60,7 @@ def open(filename, mode="rb", compresslevel=gzip_ng._COMPRESS_LEVEL_TRADEOFF,
         gzip_file = io.BufferedReader(
             _ThreadedGzipReader(filename, block_size=block_size))
     else:
-        gzip_file = io.BufferedWriter(
+        gzip_file = FlushableBufferedWriter(
             _ThreadedGzipWriter(
                 filename,
                 mode.replace("t", "b"),
@@ -101,6 +101,7 @@ def __init__(self, filename, queue_size=2, block_size=1024 * 1024):
         self.worker = threading.Thread(target=self._decompress)
         self._closed = False
         self.running = True
+        self._calling_thread = threading.current_thread()
         self.worker.start()
 
     def _check_closed(self, msg=None):
@@ -110,15 +111,15 @@ def _check_closed(self, msg=None):
     def _decompress(self):
         block_size = self.block_size
         block_queue = self.queue
-        while self.running:
+        while self.running and self._calling_thread.is_alive():
             try:
                 data = self.fileobj.read(block_size)
             except Exception as e:
                 self.exception = e
                 return
             if not data:
                 return
-            while self.running:
+            while self.running and self._calling_thread.is_alive():
                 try:
                     block_queue.put(data, timeout=0.05)
                     break
@@ -166,6 +167,12 @@ def closed(self) -> bool:
         return self._closed
 
 
+class FlushableBufferedWriter(io.BufferedWriter):
+    def flush(self):
+        super().flush()
+        self.raw.flush()
+
+
 class _ThreadedGzipWriter(io.RawIOBase):
     """
     Write a gzip file using multiple threads.
@@ -215,6 +222,7 @@ def __init__(self,
         if "b" not in mode:
             mode += "b"
         self.lock = threading.Lock()
+        self._calling_thread = threading.current_thread()
         self.exception: Optional[Exception] = None
         self.level = level
         self.previous_block = b""
@@ -313,30 +321,35 @@ def write(self, b) -> int:
         self.input_queues[worker_index].put((data, zdict))
         return len(data)
 
-    def flush(self):
+    def _end_gzip_stream(self):
         self._check_closed()
         # Wait for all data to be compressed
         for in_q in self.input_queues:
             in_q.join()
         # Wait for all data to be written
         for out_q in self.output_queues:
             out_q.join()
+        # Write an empty deflate block with a lost block marker.
+        self.raw.write(zlib_ng.compress(b"", wbits=-15))
+        trailer = struct.pack("<II", self._crc, self._size & 0xFFFFFFFF)
+        self.raw.write(trailer)
+        self._crc = 0
+        self._size = 0
         self.raw.flush()
 
+    def flush(self):
+        self._end_gzip_stream()
+        self._write_gzip_header()
+
     def close(self) -> None:
         if self._closed:
             return
-        self.flush()
+        self._end_gzip_stream()
         self.stop()
         if self.exception:
             self.raw.close()
             self._closed = True
             raise self.exception
-        # Write an empty deflate block with a lost block marker.
-        self.raw.write(zlib_ng.compress(b"", wbits=-15))
-        trailer = struct.pack("<II", self._crc, self._size & 0xFFFFFFFF)
-        self.raw.write(trailer)
-        self.raw.flush()
         if self.closefd:
             self.raw.close()
         self._closed = True
@@ -353,7 +366,7 @@ def _compress(self, index: int):
             try:
                 data, zdict = in_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
+                if not (self.running and self._calling_thread.is_alive()):
                     return
                 continue
             try:
@@ -369,41 +382,31 @@ def _compress(self, index: int):
     def _write(self):
         index = 0
         output_queues = self.output_queues
-        fp = self.raw
-        total_crc = 0
-        size = 0
         while True:
             out_index = index % self.threads
             output_queue = output_queues[out_index]
             try:
                 compressed, crc, data_length = output_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
-                    self._crc = total_crc
-                    self._size = size
+                if not (self.running and self._calling_thread.is_alive()):
                     return
                 continue
-            total_crc = zlib_ng.crc32_combine(total_crc, crc, data_length)
-            size += data_length
-            fp.write(compressed)
+            self._crc = zlib_ng.crc32_combine(self._crc, crc, data_length)
+            self._size += data_length
+            self.raw.write(compressed)
             output_queue.task_done()
             index += 1
 
     def _compress_and_write(self):
         if not self.threads == 1:
             raise SystemError("Compress_and_write is for one thread only")
-        fp = self.raw
-        total_crc = 0
-        size = 0
         in_queue = self.input_queues[0]
         compressor = self.compressors[0]
         while True:
             try:
                 data, zdict = in_queue.get(timeout=0.05)
             except queue.Empty:
-                if not self.running:
-                    self._crc = total_crc
-                    self._size = size
+                if not (self.running and self._calling_thread.is_alive()):
                     return
                 continue
             try:
@@ -413,9 +416,9 @@ def _compress_and_write(self):
                 self._set_error_and_empty_queue(e, in_queue)
                 return
             data_length = len(data)
-            total_crc = zlib_ng.crc32_combine(total_crc, crc, data_length)
-            size += data_length
-            fp.write(compressed)
+            self._crc = zlib_ng.crc32_combine(self._crc, crc, data_length)
+            self._size += data_length
+            self.raw.write(compressed)
             in_queue.task_done()
 
     def _set_error_and_empty_queue(self, error, q):

tests/test_gzip_ng_threaded.py

@@ -9,6 +9,8 @@
 import io
 import itertools
 import os
+import subprocess
+import sys
 import tempfile
 from pathlib import Path
 
@@ -209,3 +211,38 @@ def test_threaded_writer_does_not_close_stream():
     assert not test_stream.closed
     test_stream.seek(0)
     assert gzip.decompress(test_stream.read()) == b"thisisatest"
+
+
+@pytest.mark.timeout(5)
+@pytest.mark.parametrize(
+    ["mode", "threads"], itertools.product(["rb", "wb"], [1, 2]))
+def test_threaded_program_can_exit_on_error(tmp_path, mode, threads):
+    program = tmp_path / "no_context_manager.py"
+    test_file = tmp_path / "output.gz"
+    # Write 40 mb input data to saturate read buffer. Because of the repetitive
+    # nature the resulting gzip file is very small (~40 KiB).
+    test_file.write_bytes(gzip.compress(b"test" * (10 * 1024 * 1024)))
+    with open(program, "wt") as f:
+        f.write("from zlib_ng import gzip_ng_threaded\n")
+        f.write(
+            f"f = gzip_ng_threaded.open('{test_file}', "
+            f"mode='{mode}', threads={threads})\n"
+        )
+        f.write("raise Exception('Error')\n")
+    subprocess.run([sys.executable, str(program)])
+
+
+@pytest.mark.parametrize("threads", [1, 2])
+def test_flush(tmp_path, threads):
+    test_file = tmp_path / "output.gz"
+    with gzip_ng_threaded.open(test_file, "wb", threads=threads) as f:
+        f.write(b"1")
+        f.flush()
+        assert gzip.decompress(test_file.read_bytes()) == b"1"
+        f.write(b"2")
+        f.flush()
+        assert gzip.decompress(test_file.read_bytes()) == b"12"
+        f.write(b"3")
+        f.flush()
+        assert gzip.decompress(test_file.read_bytes()) == b"123"
+    assert gzip.decompress(test_file.read_bytes()) == b"123"