feat: stream torrentzip members, detect zlib-ng

write_torrentzip reads and deflates one member at a time, gives back what it read (size, crc32, sha1) and takes a ZIP64 field only past 32 bits, so a 2 GiB disk image builds without being held whole. TorrentZip layouts are unchanged. zlib_ng_version() names the deflater that cannot reproduce TorrentZip bytes.

scripts/torrentzip.py

@@ -24,15 +24,126 @@ from __future__ import annotations
 
 import argparse
 import hashlib
+import io
 import struct
 import zipfile
 import zlib
+from collections.abc import Callable, Iterable
 from pathlib import Path
+from typing import BinaryIO, NamedTuple
 
 # TorrentZip constants: 1996-12-24 23:32:00, deflate level 9
 _DOS_DATE = ((1996 - 1980) << 9) | (12 << 5) | 24
 _DOS_TIME = (23 << 11) | (32 << 5) | 0
 _COMMENT_PREFIX = "TORRENTZIPPED-"
+# Members are read and deflated this much at a time. Deflate output does not
+# depend on how its input is cut, so a 2 GiB disk image streams through in
+# the same bytes a one-shot compression gives.
+CHUNK = 16 * 1024 * 1024
+# A size or offset that does not fit 32 bits takes a ZIP64 field. No
+# TorrentZip-conformant archive needs one, so their layout is untouched.
+ZIP32_LIMIT = 0xFFFFFFFF
+_ZIP64_MARK = 0xFFFFFFFF
+
+
+class Written(NamedTuple):
+    """A member as the writer read it from its source."""
+
+    name: str
+    size: int
+    crc32: str
+    sha1: str
+
+
+def zlib_ng_version(module: object = zlib) -> str | None:
+    """The zlib-ng version this interpreter deflates with, None for zlib.
+
+    TorrentZip bytes are the output of zlib's level-9 deflate. Stock zlib
+    1.2.13 to 1.3.2 reproduces every TorrentZip archive in the collection.
+    zlib-ng 2.2.4 (python.org 3.14 on Windows) reproduces 479 of 1,561 and
+    2.3.3 545, and the two disagree with each other.
+    """
+    version = getattr(module, "ZLIBNG_VERSION", None)
+    if version:
+        return str(version)
+    runtime = str(getattr(module, "ZLIB_RUNTIME_VERSION", "") or "")
+    return runtime if "zlib-ng" in runtime.lower() else None
+
+
+def _deflate(stream: BinaryIO) -> tuple[bytes, Written]:
+    compressor = zlib.compressobj(9, zlib.DEFLATED, -15)
+    blob = bytearray()
+    crc, size, sha1 = 0, 0, hashlib.sha1()
+    while chunk := stream.read(CHUNK):
+        crc = zlib.crc32(chunk, crc)
+        size += len(chunk)
+        sha1.update(chunk)
+        blob += compressor.compress(chunk)
+    blob += compressor.flush()
+    return bytes(blob), Written("", size, f"{crc & 0xFFFFFFFF:08x}", sha1.hexdigest())
+
+
+def write_torrentzip(
+    sink: BinaryIO, members: Iterable[tuple[str, Callable[[], BinaryIO]]]
+) -> list[Written]:
+    """Write a TorrentZip archive to *sink*, one member in memory at a time.
+
+    *members* pairs each name with a callable opening its bytes. Members are
+    sorted case-insensitively by name, as the format requires. Returns what
+    was read, in archive order, so a caller can check every member.
+    """
+    ordered = sorted(members, key=lambda m: m[0].lower())
+    central = bytearray()
+    offset = 0
+    written: list[Written] = []
+    for name, opener in ordered:
+        raw = name.encode("ascii")
+        with opener() as stream:
+            blob, read = _deflate(stream)
+        crc = int(read.crc32, 16)
+        written.append(read._replace(name=name))
+        wide = read.size >= ZIP32_LIMIT or len(blob) >= ZIP32_LIMIT
+        far = offset >= ZIP32_LIMIT
+        local_extra = struct.pack("<HHQQ", 1, 16, read.size, len(blob)) if wide else b""
+        central_extra = (
+            struct.pack("<QQ", read.size, len(blob)) if wide else b""
+        ) + (struct.pack("<Q", offset) if far else b"")
+        if central_extra:
+            central_extra = struct.pack("<HH", 1, len(central_extra)) + central_extra
+        sizes = (_ZIP64_MARK, _ZIP64_MARK) if wide else (len(blob), read.size)
+        needed = 45 if wide or far else 20
+        header = struct.pack(
+            "<IHHHHHIIIHH", 0x04034B50, needed, 2, 8, _DOS_TIME, _DOS_DATE,
+            crc, *sizes, len(raw), len(local_extra),
+        ) + raw + local_extra
+        sink.write(header)
+        sink.write(blob)
+        central += struct.pack(
+            "<IHHHHHHIIIHHHHHII", 0x02014B50, 0, needed, 2, 8, _DOS_TIME, _DOS_DATE,
+            crc, *sizes, len(raw), len(central_extra), 0, 0, 0, 0,
+            _ZIP64_MARK if far else offset,
+        ) + raw + central_extra
+        offset += len(header) + len(blob)
+
+    cd_offset = offset
+    sink.write(central)
+    count = len(ordered)
+    if cd_offset >= ZIP32_LIMIT or len(central) >= ZIP32_LIMIT or count >= 0xFFFF:
+        end64 = cd_offset + len(central)
+        sink.write(struct.pack(
+            "<IQHHIIQQQQ", 0x06064B50, 44, 45, 45, 0, 0, count, count,
+            len(central), cd_offset,
+        ))
+        sink.write(struct.pack("<IIQI", 0x07064B50, 0, end64, 1))
+    comment = (
+        f"{_COMMENT_PREFIX}{zlib.crc32(bytes(central)) & 0xFFFFFFFF:08X}"
+    ).encode("ascii")
+    sink.write(struct.pack(
+        "<IHHHHIIH", 0x06054B50, 0, 0, min(count, 0xFFFF), min(count, 0xFFFF),
+        min(len(central), _ZIP64_MARK), min(cd_offset, _ZIP64_MARK), len(comment),
+    ))
+    sink.write(comment)
+    return written
 
 
 def build_torrentzip(members: list[tuple[str, bytes]]) -> bytes:
@@ -41,38 +152,11 @@ def build_torrentzip(members: list[tuple[str, bytes]]) -> bytes:
     Members are sorted case-insensitively by name, as the format requires.
     Returns the complete archive bytes.
     """
-    ordered = sorted(members, key=lambda m: m[0].lower())
-    body = bytearray()
-    central = bytearray()
-
-    for name, data in ordered:
-        raw = name.encode("ascii")
-        compressor = zlib.compressobj(9, zlib.DEFLATED, -15)
-        blob = compressor.compress(data) + compressor.flush()
-        crc = zlib.crc32(data) & 0xFFFFFFFF
-        offset = len(body)
-        body += struct.pack(
-            "<IHHHHHIIIHH", 0x04034B50, 20, 2, 8, _DOS_TIME, _DOS_DATE,
-            crc, len(blob), len(data), len(raw), 0,
-        )
-        body += raw + blob
-        central += struct.pack(
-            "<IHHHHHHIIIHHHHHII", 0x02014B50, 0, 20, 2, 8, _DOS_TIME, _DOS_DATE,
-            crc, len(blob), len(data), len(raw), 0, 0, 0, 0, 0, offset,
-        )
-        central += raw
-
-    cd_offset = len(body)
-    body += central
-    comment = (
-        f"{_COMMENT_PREFIX}{zlib.crc32(bytes(central)) & 0xFFFFFFFF:08X}"
-    ).encode("ascii")
-    body += struct.pack(
-        "<IHHHHIIH", 0x06054B50, 0, 0, len(ordered), len(ordered),
-        len(central), cd_offset, len(comment),
+    sink = io.BytesIO()
+    write_torrentzip(
+        sink, [(name, lambda data=data: io.BytesIO(data)) for name, data in members]
     )
-    body += comment
-    return bytes(body)
+    return sink.getvalue()
 
 
 def rebuild_torrentzip(source: str | Path) -> bytes:

tests/test_torrentzip.py

@@ -155,6 +155,72 @@ def _as_file(test: unittest.TestCase, data: bytes) -> str:
     return path
 
 
+class StreamingWriter(unittest.TestCase):
+    """Members stream through in chunks, one in memory at a time."""
+
+    def _members(self) -> list[tuple[str, bytes]]:
+        import random
+
+        noise = random.Random(7).randbytes(70_000)
+        return [("b.rom", noise + bytes(50_000) + noise[:9_000]), ("A.rom", b"x" * 1_000)]
+
+    def test_chunked_input_gives_the_one_shot_bytes(self):
+        """Deflate output does not depend on how its input is cut."""
+        import io
+        import torrentzip
+
+        members = self._members()
+        whole = torrentzip.build_torrentzip(members)
+        saved = torrentzip.CHUNK
+        try:
+            for chunk in (1, 333, 4096):
+                torrentzip.CHUNK = chunk
+                self.assertEqual(torrentzip.build_torrentzip(members), whole, chunk)
+        finally:
+            torrentzip.CHUNK = saved
+        sink = io.BytesIO()
+        written = torrentzip.write_torrentzip(
+            sink, [(name, lambda d=data: io.BytesIO(d)) for name, data in members]
+        )
+        self.assertEqual(sink.getvalue(), whole)
+        self.assertEqual(
+            [(w.name, w.size, w.sha1) for w in written],
+            [(name, len(data), hashlib.sha1(data).hexdigest())
+             for name, data in sorted(members, key=lambda m: m[0].lower())],
+        )
+
+    def test_a_member_past_32_bits_takes_a_zip64_field(self):
+        """cspect-next-8gb.zip holds an 8 GiB image: its size needs ZIP64."""
+        import torrentzip
+
+        members = self._members()
+        saved = torrentzip.ZIP32_LIMIT
+        torrentzip.ZIP32_LIMIT = 60_000
+        try:
+            data = torrentzip.build_torrentzip(members)
+        finally:
+            torrentzip.ZIP32_LIMIT = saved
+        with zipfile.ZipFile(_as_file(self, data)) as zf:
+            self.assertIsNone(zf.testzip())
+            self.assertEqual(
+                {info.filename: zf.read(info) for info in zf.infolist()}, dict(members)
+            )
+        self.assertNotEqual(data, build_torrentzip(members))
+
+    def test_zlib_ng_is_recognised(self):
+        """python.org 3.14 on Windows links zlib-ng, which gives other bytes."""
+        from types import SimpleNamespace
+
+        from torrentzip import zlib_ng_version
+
+        stock = SimpleNamespace(ZLIB_RUNTIME_VERSION="1.3.1")
+        windows = SimpleNamespace(ZLIB_RUNTIME_VERSION="1.3.1.zlib-ng")
+        binding = SimpleNamespace(ZLIB_RUNTIME_VERSION="1.3.1", ZLIBNG_VERSION="2.2.4")
+        self.assertIsNone(zlib_ng_version(stock))
+        self.assertEqual(zlib_ng_version(windows), "1.3.1.zlib-ng")
+        self.assertEqual(zlib_ng_version(binding), "2.2.4")
+
+
 class CheckModeExitCode(unittest.TestCase):
     """--check printed `conform: False` and exited 0, so no script could use it."""