From 4d956fa13b52eb0d7a0eea69040af91336ddfa1b Mon Sep 17 00:00:00 2001 From: stumpylog <797416+stumpylog@users.noreply.github.com> Date: Thu, 23 Jul 2026 13:27:06 -0700 Subject: [PATCH] Feature: add export compression policy module --- src/documents/export/compression.py | 101 ++++++++++++++++++ .../tests/export/test_compression.py | 98 +++++++++++++++++ 2 files changed, 199 insertions(+) create mode 100644 src/documents/export/compression.py create mode 100644 src/documents/tests/export/test_compression.py diff --git a/src/documents/export/compression.py b/src/documents/export/compression.py new file mode 100644 index 000000000..4760a16b1 --- /dev/null +++ b/src/documents/export/compression.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import importlib +import zipfile + +# ZIP_ZSTANDARD exists only on Python 3.14+ (PEP 784). None elsewhere. +ZSTD: int | None = getattr(zipfile, "ZIP_ZSTANDARD", None) + +# CLI choices are fixed across runtimes so argparse never hides zstd; runtime +# availability is enforced separately in compression_available(). +COMPRESSION_CHOICES: tuple[str, ...] = ( + "stored", + "deflated", + "bzip2", + "lzma", + "zstd", +) + +# Method name -> zipfile compression constant (zstd only when supported). +COMPRESSION_METHODS: dict[str, int] = { + "stored": zipfile.ZIP_STORED, + "deflated": zipfile.ZIP_DEFLATED, + "bzip2": zipfile.ZIP_BZIP2, + "lzma": zipfile.ZIP_LZMA, +} +if ZSTD is not None: + COMPRESSION_METHODS["zstd"] = ZSTD + +# Inclusive (min, max) level bounds per method; None => level not applicable. +# Verified on CPython 3.14.3. +LEVEL_BOUNDS: dict[str, tuple[int, int] | None] = { + "stored": None, + "deflated": (0, 9), + "bzip2": (1, 9), + "lzma": None, + "zstd": (-131072, 22), +} + +# zipfile compress_type id -> method name. 93 = current zstd id, 20 = legacy +# zstd id that zipfile can still read. +_COMPRESS_TYPE_TO_METHOD: dict[int, str] = { + zipfile.ZIP_STORED: "stored", + zipfile.ZIP_DEFLATED: "deflated", + zipfile.ZIP_BZIP2: "bzip2", + zipfile.ZIP_LZMA: "lzma", + 93: "zstd", + 20: "zstd", +} + + +def compression_available(method: str) -> bool: + """Whether the running interpreter can actually use the given method.""" + if method in ("stored", "deflated"): + # zlib is a hard CPython dependency; stored needs nothing. + return True + if method == "bzip2": + return _module_importable("bz2") + if method == "lzma": + return _module_importable("lzma") + if method == "zstd": + return ZSTD is not None and _module_importable("compression.zstd") + return False + + +def _module_importable(name: str) -> bool: + try: + importlib.import_module(name) + except ImportError: + return False + return True + + +def level_error(method: str, level: int | None) -> str | None: + """Return a human message if (method, level) is invalid, else None.""" + if level is None: + return None + bounds = LEVEL_BOUNDS[method] + if bounds is None: + return f"--zip-compression-level has no effect for '{method}'" + low, high = bounds + if not (low <= level <= high): + return ( + f"--zip-compression-level for '{method}' must be between {low} and {high}" + ) + return None + + +def compress_type_readable(compress_type: int) -> bool: + """Whether this interpreter can decompress an entry of the given type.""" + method = _COMPRESS_TYPE_TO_METHOD.get(compress_type) + if method is None: + return False + return compression_available(method) + + +def unreadable_method_names(compress_types: set[int]) -> set[str]: + """Map a set of compress_type ids to human method names for error messages.""" + names: set[str] = set() + for ct in compress_types: + names.add(_COMPRESS_TYPE_TO_METHOD.get(ct, f"method {ct}")) + return names diff --git a/src/documents/tests/export/test_compression.py b/src/documents/tests/export/test_compression.py new file mode 100644 index 000000000..e7288cff3 --- /dev/null +++ b/src/documents/tests/export/test_compression.py @@ -0,0 +1,98 @@ +import sys +import zipfile + +import pytest + +from documents.export import compression + + +class TestCompressionMethods: + def test_choices_always_include_zstd(self) -> None: + # zstd is offered regardless of runtime; availability is checked separately + assert compression.COMPRESSION_CHOICES == ( + "stored", + "deflated", + "bzip2", + "lzma", + "zstd", + ) + + @pytest.mark.parametrize( + ("name", "constant"), + [ + ("stored", zipfile.ZIP_STORED), + ("deflated", zipfile.ZIP_DEFLATED), + ("bzip2", zipfile.ZIP_BZIP2), + ("lzma", zipfile.ZIP_LZMA), + ], + ) + def test_method_maps_to_zipfile_constant(self, name: str, constant: int) -> None: + assert compression.COMPRESSION_METHODS[name] == constant + + def test_stored_and_deflated_always_available(self) -> None: + assert compression.compression_available("stored") + assert compression.compression_available("deflated") + + def test_zstd_availability_tracks_runtime(self) -> None: + expected: bool = sys.version_info >= (3, 14) + assert compression.compression_available("zstd") == expected + + +class TestLevelError: + @pytest.mark.parametrize( + ("method", "level"), + [ + ("deflated", 0), + ("deflated", 9), + ("bzip2", 1), + ("bzip2", 9), + ("deflated", None), + ("stored", None), + ], + ) + def test_valid_levels_return_none(self, method: str, level: int | None) -> None: + assert compression.level_error(method, level) is None + + @pytest.mark.parametrize( + ("method", "level"), + [ + ("deflated", 10), + ("deflated", -1), + ("bzip2", 0), + ("bzip2", 10), + ], + ) + def test_out_of_range_levels_return_message( + self, + method: str, + level: int, + ) -> None: + msg: str | None = compression.level_error(method, level) + assert msg is not None + assert "between" in msg + + @pytest.mark.parametrize("method", ["stored", "lzma"]) + def test_level_on_levelless_method_is_rejected(self, method: str) -> None: + msg: str | None = compression.level_error(method, 5) + assert msg is not None + assert "no effect" in msg + + +class TestCompressTypeReadable: + @pytest.mark.parametrize("ct", [zipfile.ZIP_STORED, zipfile.ZIP_DEFLATED]) + def test_stored_and_deflated_always_readable(self, ct: int) -> None: + assert compression.compress_type_readable(ct) + + def test_zstd_compress_type_readability_tracks_runtime(self) -> None: + # 93 = ZIP_ZSTANDARD; 20 = legacy zstd method id (read-only) + expected: bool = sys.version_info >= (3, 14) + assert compression.compress_type_readable(93) == expected + assert compression.compress_type_readable(20) == expected + + def test_unknown_compress_type_is_unreadable(self) -> None: + assert not compression.compress_type_readable(9999) + + def test_unreadable_method_names_lists_methods(self) -> None: + # An unknown method id maps to no name and is reported generically. + names: set[str] = compression.unreadable_method_names({9999}) + assert names == {"method 9999"}