Add opt-in lossless TIFF compression writer settings

src/polystore/config.py

@@ -3,14 +3,18 @@
 from __future__ import annotations
 
 from abc import ABC, abstractmethod
+from collections.abc import Sequence
 from dataclasses import dataclass
 from enum import Enum
+from pathlib import Path
 from typing import ClassVar
 
 from metaclass_registry import AutoRegisterMeta
-from numcodecs import Blosc, LZ4, Zlib, Zstd
+from numcodecs import LZ4, Blosc, Zlib, Zstd
 from numcodecs.abc import Codec
 
+from .formats import FileFormat
+
 
 class ZarrCompressor(Enum):
     """Compression algorithms supported by :class:`ZarrStorageBackend`."""
@@ -116,6 +120,84 @@ class ZarrChunkStrategy(Enum):
     FILE = "file"
 
 
+class TiffCompression(Enum):
+    """Lossless compression choices for disk-backed TIFF output."""
+
+    NONE = ("none", None)
+    DEFLATE = ("deflate", "zlib")
+
+    def __new__(cls, value: str, tifffile_codec: str | None):
+        member = object.__new__(cls)
+        member._value_ = value
+        member.tifffile_codec = tifffile_codec
+        return member
+
+
+@dataclass(frozen=True)
+class TiffConfig:
+    """TIFF writer settings; defaults preserve the uncompressed file contract.
+
+    Args:
+        compression: Lossless disk-TIFF codec. ``NONE`` preserves the existing
+            uncompressed output; ``DEFLATE`` compresses TIFF pixels exactly.
+        compression_level: Deflate level from 1 to 9. Ignored when compression
+            is ``NONE``.
+    """
+
+    compression: TiffCompression = TiffCompression.NONE
+    compression_level: int = 1
+
+    def __post_init__(self) -> None:
+        if not isinstance(self.compression, TiffCompression):
+            raise TypeError("compression must be a TiffCompression value")
+        if isinstance(self.compression_level, bool) or not isinstance(
+            self.compression_level, int
+        ) or not 1 <= self.compression_level <= 9:
+            raise ValueError("compression_level must be an integer from 1 to 9")
+
+    def tifffile_write_kwargs(self) -> dict[str, object]:
+        """Return only codec options accepted by tifffile.imwrite."""
+
+        codec = self.compression.tifffile_codec
+        if codec is None:
+            return {}
+        return {
+            "compression": codec,
+            "compressionargs": {"level": self.compression_level},
+        }
+
+    def applies_to_path(self, path: str | Path) -> bool:
+        """Whether this configured codec applies to this disk output path."""
+
+        return (
+            self.compression is not TiffCompression.NONE
+            and Path(path).suffix.lower() in FileFormat.TIFF.extensions
+        )
+
+
+def tiff_write_batches(
+    paths: Sequence[str | Path],
+    config: TiffConfig | None,
+) -> tuple[tuple[tuple[int, ...], TiffConfig | None], ...]:
+    """Partition consecutive writes by their exact TIFF codec requirement."""
+
+    if not paths:
+        return ()
+    if config is None or config.compression is TiffCompression.NONE:
+        return ((tuple(range(len(paths))), None),)
+    batches: list[tuple[tuple[int, ...], TiffConfig | None]] = []
+    start = 0
+    active = config.applies_to_path(paths[0])
+    for index, path in enumerate(paths[1:], start=1):
+        applies = config.applies_to_path(path)
+        if applies != active:
+            batches.append((tuple(range(start, index)), config if active else None))
+            start = index
+            active = applies
+    batches.append((tuple(range(start, len(paths))), config if active else None))
+    return tuple(batches)
+
+
 @dataclass(frozen=True)
 class ZarrConfig:
     """Framework-independent configuration for Zarr storage.

src/polystore/disk.py

@@ -18,6 +18,7 @@
 
 from .array_payload import storage_numpy_array
 from .base import StorageBackend
+from .config import TiffConfig
 from .constants import Backend
 from .exceptions import StorageResolutionError
 from .formats import FileFormat
@@ -199,9 +200,18 @@ def _tensorflow_reader(self, path, **kwargs):
         tf = FileFormat.TENSORFLOW.load_dependency()
         return tf.io.parse_tensor(tf.io.read_file(path.as_posix()), out_type=tf.dtypes.float32)
 
-    def _tiff_writer(self, path, data, **kwargs):
+    def _tiff_writer(self, path, data, *, tiff_config: TiffConfig | None = None, **kwargs):
+        """Write TIFF pixels with optional lossless, per-call codec settings."""
+
         tifffile = FileFormat.TIFF.load_dependency()
-        tifffile.imwrite(path, storage_numpy_array(data))
+        config = TiffConfig() if tiff_config is None else tiff_config
+        if not isinstance(config, TiffConfig):
+            raise TypeError("tiff_config must be a TiffConfig value")
+        tifffile.imwrite(
+            path,
+            storage_numpy_array(data),
+            **config.tifffile_write_kwargs(),
+        )
 
     def _tiff_reader(self, path):
         tifffile = FileFormat.TIFF.load_dependency()

tests/test_disk_backend.py

@@ -4,16 +4,17 @@
 ensure_directory idempotence, and symlink creation.
 """
 
-from pathlib import Path
 import hashlib
+from pathlib import Path
 
 import numpy as np
-import tifffile
 import pytest
+import tifffile
 
+from polystore import FileManager
 from polystore import disk as disk_module
+from polystore.config import TiffCompression, TiffConfig, tiff_write_batches
 from polystore.disk import DiskBackend
-from polystore import FileManager
 from polystore.exceptions import StorageResolutionError
 
 
@@ -109,6 +110,47 @@ class ExternalArray:
     np.testing.assert_array_equal(tifffile.imread(path), expected)
 
 
+def test_tiff_deflate_is_lossless_and_opt_in(tmp_path: Path) -> None:
+    pixels = np.zeros((2, 256, 256), dtype=np.int32)
+    pixels[0, 20:40, 30:50] = 70001
+    raw_path = tmp_path / "raw.tif"
+    compressed_path = tmp_path / "compressed.tif"
+    backend = DiskBackend()
+
+    backend.save(pixels, raw_path)
+    backend.save(
+        pixels,
+        compressed_path,
+        tiff_config=TiffConfig(
+            compression=TiffCompression.DEFLATE,
+            compression_level=3,
+        ),
+    )
+
+    with (
+        tifffile.TiffFile(raw_path) as raw,
+        tifffile.TiffFile(compressed_path) as compressed,
+    ):
+        assert raw.pages[0].compression.value == 1
+        assert compressed.pages[0].compression.value != 1
+    assert compressed_path.stat().st_size < raw_path.stat().st_size
+    np.testing.assert_array_equal(tifffile.imread(compressed_path), pixels)
+
+
+def test_tiff_write_batches_preserve_order_and_exclude_other_formats() -> None:
+    config = TiffConfig(compression=TiffCompression.DEFLATE)
+    assert tiff_write_batches(("a.csv", "b.tif", "c.tiff", "d.npy"), config) == (
+        ((0,), None),
+        ((1, 2), config),
+        ((3,), None),
+    )
+    assert tiff_write_batches(("a.csv", "b.tif"), TiffConfig()) == (
+        ((0, 1), None),
+    )
+    with pytest.raises(ValueError, match="compression_level"):
+        TiffConfig(compression_level=10)
+
+
 def test_list_files_recursive_and_extension_filter(tmp_path: Path):
     disk = DiskBackend()
     base = tmp_path / "root"