Use _GzipReader and add tests

Author: rhpvorderman

src/zlib_ng/gzip_ng.py

@@ -28,6 +28,7 @@
 import _compression  # noqa: I201  # Not third-party
 
 from . import zlib_ng
+from .zlib_ng import _GzipReader
 
 __all__ = ["GzipFile", "open", "compress", "decompress", "BadGzipFile",
            "READ_BUFFER_SIZE"]
@@ -36,19 +37,14 @@
 _COMPRESS_LEVEL_TRADEOFF = zlib_ng.Z_DEFAULT_COMPRESSION
 _COMPRESS_LEVEL_BEST = zlib_ng.Z_BEST_COMPRESSION
 
-#: The amount of data that is read in at once when decompressing a file.
-#: Increasing this value may increase performance.
-#: 128K is also the size used by pigz and cat to read files from the
-# filesystem.
-READ_BUFFER_SIZE = 128 * 1024
+# The amount of data that is read in at once when decompressing a file.
+# Increasing this value may increase performance.
+READ_BUFFER_SIZE = 512 * 1024
 
 FTEXT, FHCRC, FEXTRA, FNAME, FCOMMENT = 1, 2, 4, 8, 16
 READ, WRITE = 1, 2
 
-try:
-    BadGzipFile = gzip.BadGzipFile  # type: ignore
-except AttributeError:  # Versions lower than 3.8 do not have BadGzipFile
-    BadGzipFile = OSError  # type: ignore
+BadGzipFile = gzip.BadGzipFile  # type: ignore
 
 
 # The open method was copied from the CPython source with minor adjustments.
@@ -149,7 +145,7 @@ def __init__(self, filename=None, mode=None,
                                                 zlib_ng.DEF_MEM_LEVEL,
                                                 0)
         if self.mode == READ:
-            raw = _GzipNGReader(self.fileobj)
+            raw = _GzipReader(self.fileobj, READ_BUFFER_SIZE)
             self._buffer = io.BufferedReader(raw)
 
     def __repr__(self):
@@ -180,73 +176,9 @@ def write(self, data):
         return length
 
 
-class _GzipNGReader(gzip._GzipReader):
-    def __init__(self, fp):
-        # Call the init method of gzip._GzipReader's parent here.
-        # It is not very invasive and allows us to override _PaddedFile
-        _compression.DecompressReader.__init__(
-            self, gzip._PaddedFile(fp), zlib_ng._ZlibDecompressor,
-            wbits=-zlib_ng.MAX_WBITS)
-        # Set flag indicating start of a new member
-        self._new_member = True
-        self._last_mtime = None
-
-    def read(self, size=-1):
-        if size < 0:
-            return self.readall()
-        # size=0 is special because decompress(max_length=0) is not supported
-        if not size:
-            return b""
-
-        # For certain input data, a single
-        # call to decompress() may not return
-        # any data. In this case, retry until we get some data or reach EOF.
-        while True:
-            if self._decompressor.eof:
-                # Ending case: we've come to the end of a member in the file,
-                # so finish up this member, and read a new gzip header.
-                # Check the CRC and file size, and set the flag so we read
-                # a new member
-                self._read_eof()
-                self._new_member = True
-                self._decompressor = self._decomp_factory(
-                    **self._decomp_args)
-
-            if self._new_member:
-                # If the _new_member flag is set, we have to
-                # jump to the next member, if there is one.
-                self._init_read()
-                if not self._read_gzip_header():
-                    self._size = self._pos
-                    return b""
-                self._new_member = False
-
-            # Read a chunk of data from the file
-            if self._decompressor.needs_input:
-                buf = self._fp.read(READ_BUFFER_SIZE)
-                uncompress = self._decompressor.decompress(buf, size)
-            else:
-                uncompress = self._decompressor.decompress(b"", size)
-            if self._decompressor.unused_data != b"":
-                # Prepend the already read bytes to the fileobj so they can
-                # be seen by _read_eof() and _read_gzip_header()
-                self._fp.prepend(self._decompressor.unused_data)
-
-            if uncompress != b"":
-                break
-            if buf == b"":
-                raise EOFError("Compressed file ended before the "
-                               "end-of-stream marker was reached")
-
-        self._crc = zlib_ng.crc32(uncompress, self._crc)
-        self._stream_size += len(uncompress)
-        self._pos += len(uncompress)
-        return uncompress
-
-
 # Aliases for improved compatibility with CPython gzip module.
 GzipFile = GzipNGFile
-_GzipReader = _GzipNGReader
+_GzipNGReader = _GzipReader
 
 
 def _read_exact(fp, n):
@@ -342,25 +274,9 @@ def decompress(data):
     """Decompress a gzip compressed string in one shot.
     Return the decompressed string.
     """
-    decompressed_members = []
-    while True:
-        fp = io.BytesIO(data)
-        if _read_gzip_header(fp) is None:
-            return b"".join(decompressed_members)
-        # Use a zlib raw deflate compressor
-        do = zlib_ng.decompressobj(wbits=-zlib_ng.MAX_WBITS)
-        # Read all the data except the header
-        decompressed = do.decompress(data[fp.tell():])
-        if not do.eof or len(do.unused_data) < 8:
-            raise EOFError("Compressed file ended before the end-of-stream "
-                           "marker was reached")
-        crc, length = struct.unpack("<II", do.unused_data[:8])
-        if crc != zlib_ng.crc32(decompressed):
-            raise BadGzipFile("CRC check failed")
-        if length != (len(decompressed) & 0xffffffff):
-            raise BadGzipFile("Incorrect length of data produced")
-        decompressed_members.append(decompressed)
-        data = do.unused_data[8:].lstrip(b"\x00")
+    fp = io.BytesIO(data)
+    reader = _GzipReader(fp, max(len(data), 16))
+    return reader.readall()
 
 
 def _argument_parser():

tests/data/test.fastq.bgzip.gz

@@ -373,9 +373,76 @@ def test_truncated_header(trunc):
         gzip_ng.decompress(trunc)
 
 
+def test_very_long_header_in_data():
+    # header with a very long filename.
+    header = (b"\x1f\x8b\x08\x08\x00\x00\x00\x00\x00\xff" + 256 * 1024 * b"A" +
+              b"\x00")
+    compressed = header + zlib.compress(b"", 3, -15) + 8 * b"\00"
+    assert gzip_ng.decompress(compressed) == b""
+
+
+def test_very_long_header_in_file():
+    # header with a very long filename.
+    header = (b"\x1f\x8b\x08\x08\x00\x00\x00\x00\x00\xff" +
+              gzip_ng.READ_BUFFER_SIZE * 2 * b"A" +
+              b"\x00")
+    compressed = header + zlib.compress(b"", 3, -15) + 8 * b"\00"
+    f = io.BytesIO(compressed)
+    with gzip_ng.open(f) as gzip_file:
+        assert gzip_file.read() == b""
+
+
 def test_concatenated_gzip():
     concat = Path(__file__).parent / "data" / "concatenated.fastq.gz"
     data = gzip.decompress(concat.read_bytes())
     with gzip_ng.open(concat, "rb") as gzip_ng_h:
         result = gzip_ng_h.read()
     assert data == result
+
+
+def test_seek():
+    from io import SEEK_CUR, SEEK_END, SEEK_SET
+    with tempfile.NamedTemporaryFile("wb", delete=False) as tmpfile:
+        tmpfile.write(gzip.compress(b"X" * 500 + b"A" + b"X" * 499))
+        tmpfile.write(gzip.compress(b"X" * 500 + b"B" + b"X" * 499))
+        tmpfile.write(gzip.compress(b"X" * 500 + b"C" + b"X" * 499))
+        tmpfile.write(gzip.compress(b"X" * 500 + b"D" + b"X" * 499))
+    with gzip_ng.open(tmpfile.name, "rb") as gzip_file:
+        # Start testing forward seek
+        gzip_file.seek(500)
+        assert gzip_file.read(1) == b"A"
+        gzip_file.seek(1500)
+        assert gzip_file.read(1) == b"B"
+        # Test reverse
+        gzip_file.seek(500)
+        assert gzip_file.read(1) == b"A"
+        # Again, but with explicit SEEK_SET
+        gzip_file.seek(500, SEEK_SET)
+        assert gzip_file.read(1) == b"A"
+        gzip_file.seek(1500, SEEK_SET)
+        assert gzip_file.read(1) == b"B"
+        gzip_file.seek(500, SEEK_SET)
+        assert gzip_file.read(1) == b"A"
+        # Seeking from current position
+        gzip_file.seek(500)
+        gzip_file.seek(2000, SEEK_CUR)
+        assert gzip_file.read(1) == b"C"
+        gzip_file.seek(-1001, SEEK_CUR)
+        assert gzip_file.read(1) == b"B"
+        # Seeking from end
+        # Any positive number should end up at the end
+        gzip_file.seek(200, SEEK_END)
+        assert gzip_file.read(1) == b""
+        gzip_file.seek(-1500, SEEK_END)
+        assert gzip_file.read(1) == b"C"
+    os.remove(tmpfile.name)
+
+
+def test_bgzip():
+    bgzip_file = Path(__file__).parent / "data" / "test.fastq.bgzip.gz"
+    gzip_file = Path(__file__).parent / "data" / "test.fastq.gz"
+    with gzip_ng.open(bgzip_file, "rb") as bgz:
+        bgz_data = bgz.read()
+    with gzip_ng.open(gzip_file, "rb") as gz:
+        gz_data = gz.read()
+    assert bgz_data == gz_data