diff --git a/src/polystore/config.py b/src/polystore/config.py index 81366ac..af36d0e 100644 --- a/src/polystore/config.py +++ b/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. diff --git a/src/polystore/disk.py b/src/polystore/disk.py index bcd5f74..e31daa2 100644 --- a/src/polystore/disk.py +++ b/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() diff --git a/tests/test_disk_backend.py b/tests/test_disk_backend.py index fb39828..a452b08 100644 --- a/tests/test_disk_backend.py +++ b/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"