diff --git a/src/polystore/base.py b/src/polystore/base.py index 921c13f..442f34c 100644 --- a/src/polystore/base.py +++ b/src/polystore/base.py @@ -209,6 +209,13 @@ def supports_file_path(self, path: Union[str, Path]) -> bool: del path return self.supports_arbitrary_files + def prepare_process(self) -> None: + """Pay this backend's one-time library costs now instead of on first I/O. + + A serving process calls this before it forks workers and reports ready. + Backends whose libraries load plugins or codecs lazily override it. + """ + def bind_registry(self, registry: Mapping[str, "BackendBase"]) -> None: """Bind an execution-local backend registry when this backend needs it.""" del registry @@ -779,6 +786,13 @@ def ensure_storage_registry() -> None: logger.info("Lazily initialized storage registry") +def prepare_storage_backends() -> None: + """Build the storage registry and prepare every backend in it for first I/O.""" + ensure_storage_registry() + for backend in storage_registry.values(): + backend.prepare_process() + + def get_backend(backend_type: str) -> DataSink: """ Get a backend by type, ensuring registry is initialized. diff --git a/src/polystore/disk.py b/src/polystore/disk.py index 96a4be6..9020210 100644 --- a/src/polystore/disk.py +++ b/src/polystore/disk.py @@ -159,6 +159,39 @@ def _disk_format_registrations(self) -> tuple[DiskFileFormatRegistration, ...]: ), ) + def prepare_process(self) -> None: + """Load TIFF codecs and raster image plugins before the first image I/O.""" + self._load_tiff_codecs() + self._load_raster_image_plugins() + + @staticmethod + def _load_tiff_codecs() -> None: + """Resolve every codec tifffile can use; it imports each on first lookup.""" + if not FileFormat.TIFF.dependency.is_installed: + return + tifffile = FileFormat.TIFF.load_dependency() + codec_tables = ( + (tifffile.COMPRESSION, tifffile.TIFF.COMPRESSORS), + (tifffile.COMPRESSION, tifffile.TIFF.DECOMPRESSORS), + (tifffile.PREDICTOR, tifffile.TIFF.PREDICTORS), + (tifffile.PREDICTOR, tifffile.TIFF.UNPREDICTORS), + ) + for kinds, table in codec_tables: + for kind in kinds: + # Membership resolves and caches the codec; False means the + # codec's optional library is absent. + kind in table # noqa: B015 + + @staticmethod + def _load_raster_image_plugins() -> None: + """Round-trip one pixel through imageio; its first I/O loads the plugins.""" + imageio = FileFormat.PNG.load_dependency() + extension = FileFormat.PNG.extensions[0] + encoded = imageio.imwrite( + "", np.zeros((1, 1), dtype=np.uint8), extension=extension + ) + imageio.imread(encoded, extension=extension) + # Format-specific writer/reader functions (pickleable) # Only needed for formats that require special handling beyond library save/load diff --git a/src/polystore/filemanager.py b/src/polystore/filemanager.py index 26cd111..2a728c1 100644 --- a/src/polystore/filemanager.py +++ b/src/polystore/filemanager.py @@ -28,6 +28,7 @@ from .exceptions import StorageResolutionError from .formats import DEFAULT_IMAGE_EXTENSIONS +from .utils import natural_sort from .virtual_workspace import SourcePixelRef logger = logging.getLogger(__name__) @@ -480,9 +481,6 @@ def list_image_files( """ backend_instance = self._get_backend(backend) - # List image files and apply natural sorting - from .utils import natural_sort - files = backend_instance.list_files(str(directory), pattern, extensions, recursive) return natural_sort([str(path) for path in files]) @@ -522,9 +520,6 @@ def list_files( # Get backend instance backend_instance = self._get_backend(backend) - # List files and apply natural sorting - from .utils import natural_sort - files = backend_instance.list_files( str(directory), pattern, extensions, recursive, **kwargs ) @@ -568,9 +563,6 @@ def list_dir(self, path: str | Path, backend: str) -> list[str]: backend_instance = self._get_backend(backend) try: - # Get directory listing and apply natural sorting - from .utils import natural_sort - entries = backend_instance.list_dir(str(path)) return natural_sort(entries) except (FileNotFoundError, NotADirectoryError): @@ -975,9 +967,6 @@ def collect_dirs_and_files( ) continue - # Apply natural sorting to both dirs and files before returning - from .utils import natural_sort - return natural_sort(dirs), natural_sort(files) def is_file(self, path: str | Path, backend: str) -> bool: diff --git a/src/polystore/formats.py b/src/polystore/formats.py index b1d5814..32800f5 100644 --- a/src/polystore/formats.py +++ b/src/polystore/formats.py @@ -5,6 +5,7 @@ """ import importlib +import importlib.util from enum import Enum from types import ModuleType @@ -21,6 +22,11 @@ class FileFormatDependency(Enum): TIFFFILE = "tifffile" IMAGEIO = "imageio.v3" + @property + def is_installed(self) -> bool: + """Whether this runtime's distribution is present, without importing it.""" + return importlib.util.find_spec(self.value.partition(".")[0]) is not None + def load(self) -> ModuleType: """Import this runtime only when a format operation requires it.""" try: diff --git a/tests/test_disk_backend.py b/tests/test_disk_backend.py index acb609a..d40e80e 100644 --- a/tests/test_disk_backend.py +++ b/tests/test_disk_backend.py @@ -275,3 +275,23 @@ def test_symlink_creation_and_detection(tmp_path: Path): # create symlink disk.create_symlink(src_file, link) assert disk.is_symlink(link) + + +def test_prepare_process_loads_tiff_codecs_and_raster_plugins(): + """Preparation leaves no codec or plugin import for the first image I/O.""" + import sys + + import tifffile + from PIL import Image + + from polystore.base import prepare_storage_backends, storage_registry + + prepare_storage_backends() + + assert "disk" in dict.keys(storage_registry) + assert "imageio.plugins.pillow" in sys.modules + assert Image._initialized == 2 + loaded = set(sys.modules) + for compression in tifffile.COMPRESSION: + compression in tifffile.TIFF.DECOMPRESSORS + assert set(sys.modules) == loaded