Add tarball codec and tests

Author: lmmx

src/range_streams/codecs/__init__.py

@@ -5,6 +5,8 @@
 The currently supported list of codecs is:
 - .zip
 - .conda (zip containing zstd-compressed tarballs, used for the conda package archives)
+- .tar (uncompressed only)
+- .png
 
 There are planned extensions to other archive and image formats.
 """

src/range_streams/codecs/tar/data.py

@@ -55,6 +55,9 @@ class HeaderData(SimpleDataClass):
     _H_CHECKSUM_SIZE = 8
     _H_LINK_INDICATOR_SIZE = 1
     _H_LINKED_NAME_SIZE = 100
+    _H_PAD_SIZE = 512
+    # Standard end-of-file padding is 2 padding records
+    _H_END_PAD_SIZE = 2 * _H_PAD_SIZE
 
 
 class TarData:

src/range_streams/codecs/tar/stream.py

@@ -3,11 +3,9 @@
 import io
 import struct
 
-from pyzstd import ZstdFile
 from ranges import Range
 
 from ...range_stream import RangeStream
-from ..zstd import ZstdTarFile
 from .data import COMPRESSIONS, TarData
 
 __all__ = ["ZipStream"]
@@ -20,10 +18,10 @@ def __init__(
         client=None,
         byte_range: Range | tuple[int, int] = Range("[0, 0)"),
         pruning_level: int = 0,
-        scan_header: bool = True,
+        scan_headers: bool = True,
     ):
         """
-        As for RangeStream, but if `scan_header` is True, then immediately call
+        As for RangeStream, but if `scan_headers` is True, then immediately call
         :meth:`check_header_rec` on initialisation (which will perform the necessary
         of range request to identify the files in the tar from the header record),
         setting :attr:`tarred_files`, and :meth:`~RangeStream.add` their file content
@@ -43,17 +41,86 @@ def __init__(
             url=url, client=client, byte_range=byte_range, pruning_level=pruning_level
         )
         self.data = TarData()
-        if scan_header:
-            self.check_header_rec()
-            # self.add_file_ranges()
+        if scan_headers:
+            self.check_header_recs()
+            self.add_file_ranges()
 
-    def check_header_rec(self):
-        head_byte_range = Range(0, 257)  # rest of first 512 bytes is padding
-        self.add(head_byte_range)
-        start_bytes = self.active_range_response.read()
+    def check_header_recs(self):
+        """
+        Scan through all header records in the file, building a list of
+        :class:`range_streams.codecs.tar.TarredFileInfo` objects describing the
+        files described by the headers (but do not download those corresponding
+        archived file ranges).
+
+        For efficiency, only look at the particular fields of interest, not the
+        entire header each time.
+        """
+        self.tarred_files: list[TarredFileInfo] = []
+        scan_tell = 0
+        assert self.total_bytes is not None
+        while scan_tell < (self.total_bytes - self.data.HEADER._H_END_PAD_SIZE):
+            try:
+                file_name = self.read_file_name(start_pos_offset=scan_tell)
+            except StopIteration:
+                # Expected if a tarball has more than 2 end-of-file padding records
+                break
+            file_size = self.read_file_size(start_pos_offset=scan_tell)
+            pad_size = self.data.HEADER._H_PAD_SIZE
+            pad_remainder = file_size % pad_size
+            file_padding = (pad_size - pad_remainder) if pad_remainder else 0
+            file_end_offset = pad_size + file_size + file_padding
+            tf_info = TarredFileInfo(
+                size=file_size,
+                padded_size=file_end_offset,
+                filename_length=len(file_name),
+                header_offset=scan_tell,
+                filename=file_name,
+            )
+            self.tarred_files.append(tf_info)
+            scan_tell += (
+                file_end_offset  # increment to move the cursor to the next file
+            )
+
+    def read_file_name(self, start_pos_offset: int = 0) -> str:
+        file_name_rng_start = start_pos_offset + self.data.HEADER._H_FILENAME_START
+        file_name_rng_end = file_name_rng_start + self.data.HEADER._H_FILENAME_SIZE
+        file_name_rng = Range(file_name_rng_start, file_name_rng_end)
+        self.add(file_name_rng)
+        file_name_b = self.active_range_response.read().rstrip(b"\x00")
+        if file_name_b == b"":
+            raise StopIteration("Expected file name, got padding bytes")
+        return file_name_b.decode("ascii")
+
+    def read_file_size(self, start_pos_offset: int = 0) -> int:
+        file_size_rng_start = start_pos_offset + self.data.HEADER._H_FILE_SIZE_START
+        file_size_rng_end = file_size_rng_start + self.data.HEADER._H_FILE_SIZE_SIZE
+        file_size_rng = Range(file_size_rng_start, file_size_rng_end)
+        self.add(file_size_rng)
+        file_size_b = self.active_range_response.read()
+        file_size = int(file_size_b, 8)  # convert octal number from bitstring
+        return file_size
+
+    def add_file_ranges(self):
+        for tf_info in self.tarred_files:
+            assert tf_info.filename is not None
+            self.add(tf_info.file_range, name=tf_info.filename)
+
+    @property
+    def filename_list(self) -> list[str]:
+        """
+        Return the names of files stored in
+        :attr:`~range_streams.codecs.tar.TarStream.tarred_files`.
+        """
+        if not hasattr(self, "tarred_files"):  # pragma: no cover
+            self.check_header_recs()
+        return [f.filename for f in self.tarred_files if f.filename is not None]
 
 
 class HeaderInfo:
+    """
+    Not used, may be useful if extending the class. Note USTAR format variant.
+    """
+
     _H_FILENAME = 0
     _H_FILE_MODE = 1
     _H_OWNER_UID = 2
@@ -75,71 +142,27 @@ class TarredFileInfo(HeaderInfo):
 
     def __init__(
         self,
-        # signature: bytes | int,
-        # flags: bytes | int,
-        # compress_type: bytes | int,
-        # compressed_size: bytes | int,
-        # uncompressed_size: bytes | int,
-        # filename_length: bytes | int,
-        # extra_field_length: bytes | int,
-        # comment_length: bytes | int,
-        # local_header_offset: bytes | int,
-        # filename: str | None,
+        size: int,  # ignoring header and trailing padding
+        padded_size: bytes | int,  # including both header and trailing padding
+        filename_length: bytes | int,
+        header_offset: int,
+        filename: str | None,
     ):
-        pass
-        # self.signature = signature
-        # self.flags = flags
-        # self.compress_type = compress_type
-        # self.compressed_size = compressed_size
-        # self.uncompressed_size = uncompressed_size
-        # self.filename_length = filename_length
-        # self.extra_field_length = extra_field_length
-        # self.comment_length = comment_length
-        # self.local_header_offset = local_header_offset
-        # self.filename = filename
+        self.size = size
+        self.padded_size = padded_size
+        self.filename_length = filename_length
+        self.header_offset = header_offset
+        self.filename = filename
 
     def __repr__(self):
         return (
             f"{self.__class__.__name__}"
-            # f" '{self.filename if self.filename is not None else ''}'"
-            # f" @ {self.local_header_offset!r}: {self.compressed_size!r}B"
+            f" '{self.filename if self.filename is not None else ''}'"
+            f" @ {self.header_offset!r}: {self.size!r}B"
         )
 
-    # @classmethod
-    # def from_central_directory_entry(
-    #    cls,
-    #    cd_entry: tuple,
-    #    filename: str | None = None,
-    # ):
-    #    """
-    #    Instantiate directly from an unpacked central directory struct
-    #    (describing the zipped file entry in a standardised entry order).
-    #    """
-    #    signature = cd_entry[cls._CD_SIGNATURE]
-    #    flags = cd_entry[cls._CD_FLAG_BITS]
-    #    compress_type = cd_entry[cls._CD_COMPRESS_TYPE]
-    #    compressed_size = cd_entry[cls._CD_COMPRESSED_SIZE]
-    #    uncompressed_size = cd_entry[cls._CD_UNCOMPRESSED_SIZE]
-    #    filename_length = cd_entry[cls._CD_FILENAME_LENGTH]
-    #    extra_field_length = cd_entry[cls._CD_EXTRA_FIELD_LENGTH]
-    #    comment_length = cd_entry[cls._CD_COMMENT_LENGTH]
-    #    local_header_offset = cd_entry[cls._CD_LOCAL_HEADER_OFFSET]
-    #    return cls(
-    #        signature=signature,
-    #        flags=flags,
-    #        compress_type=compress_type,
-    #        compressed_size=compressed_size,
-    #        uncompressed_size=uncompressed_size,
-    #        filename_length=filename_length,
-    #        extra_field_length=extra_field_length,
-    #        comment_length=comment_length,
-    #        local_header_offset=local_header_offset,
-    #        filename=filename,
-    #    )
-
-    # @property
-    # def file_range(self):
-    #    sig_start = self.local_header_offset
-    #    start = sig_start + ZipData().LOC_F_H.get_size() + self.filename_length
-    #    end = start + self.compressed_size
-    #    return Range(start, end)
+    @property
+    def file_range(self):
+        start = self.header_offset
+        end = start + self.size
+        return Range(start, end)

tests/codecs/data.py

@@ -1,6 +1,7 @@
 __all__ = [
     "EXAMPLE_ZIP_URL",
     "EXAMPLE_CONDA_URL",
+    "EXAMPLE_TAR_URL",
     "EXAMPLE_PNG_URL",
     "EXAMPLE_SEMITRANSPARENT_PNG_URL",
 ]
@@ -9,5 +10,6 @@
 
 EXAMPLE_ZIP_URL = f"{data_dir_URL}example_text_file.txt.zip"
 EXAMPLE_CONDA_URL = f"{data_dir_URL}tqdm-4.61.1-pyhd3eb1b0_1.conda"
+EXAMPLE_TAR_URL = f"{data_dir_URL}data.tar"
 EXAMPLE_PNG_URL = f"{data_dir_URL}red_square.png"
 EXAMPLE_SEMITRANSPARENT_PNG_URL = f"{data_dir_URL}red_square_rgba_semitransparent.png"

tests/codecs/tar_test.py

@@ -0,0 +1,56 @@
+from __future__ import annotations
+
+from pytest import fixture, mark, raises
+from ranges import Range
+
+from range_streams.codecs import TarStream
+
+from .data import EXAMPLE_TAR_URL
+
+
+@fixture(scope="session")
+def example_tar_stream():
+    return TarStream(url=EXAMPLE_TAR_URL)
+
+
+@mark.parametrize("expected", [8192])
+def test_tar_total_bytes(example_tar_stream, expected):
+    assert example_tar_stream.total_bytes == expected
+
+
+@mark.parametrize(
+    "expected", [(["red_square_rgba_semitransparent.png", "example_text_file.txt"])]
+)
+def test_tar_list_files(example_tar_stream, expected):
+    assert example_tar_stream.filename_list == expected
+
+
+@mark.parametrize(
+    "file_i,size,padded_size,fname,fname_len,header_offset",
+    [
+        (0, 5124, 6144, "red_square_rgba_semitransparent.png", 35, 0),
+        (1, 11, 1024, "example_text_file.txt", 21, 6144),
+    ],
+)
+def test_tarred_file_contents(
+    example_tar_stream, file_i, size, padded_size, fname, fname_len, header_offset
+):
+    tf_l = example_tar_stream.tarred_files
+    assert len(tf_l) == 2
+    tf = tf_l[file_i]
+    assert tf.size == size
+    assert tf.padded_size == padded_size
+    assert tf.filename == fname
+    assert tf.filename_length == fname_len
+    assert tf.header_offset == header_offset
+
+
+@mark.parametrize(
+    "file_i,expected",
+    [
+        (0, "TarredFileInfo 'red_square_rgba_semitransparent.png' @ 0: 5124B"),
+        (1, "TarredFileInfo 'example_text_file.txt' @ 6144: 11B"),
+    ],
+)
+def test_tar_repr(example_tar_stream, file_i, expected):
+    assert example_tar_stream.tarred_files[file_i].__repr__() == expected