polystore 0.1.6__py3-none-any.whl

polystore/__init__.py

@@ -0,0 +1,138 @@
+"""
+Polystore package exports.
+"""
+
+__version__ = "0.1.6"
+
+import os
+
+from .atomic import file_lock, atomic_write_json, atomic_update_json, FileLockError, FileLockTimeoutError
+from .backend_registry import (
+    get_backend_instance,
+    cleanup_backend_connections,
+    cleanup_all_backends,
+    register_cleanup_callback,
+    STORAGE_BACKENDS,
+)
+from .base import (
+    BackendBase,
+    DataSink,
+    DataSource,
+    ReadOnlyBackend,
+    StorageBackend,
+    storage_registry,
+    reset_memory_backend,
+    ensure_storage_registry,
+    get_backend,
+)
+from .constants import Backend, MemoryType, TransportMode
+from .disk import DiskStorageBackend
+from .filemanager import FileManager
+from .formats import FileFormat, DEFAULT_IMAGE_EXTENSIONS
+from .memory import MemoryStorageBackend
+from .metadata_writer import (
+    AtomicMetadataWriter,
+    MetadataWriteError,
+    METADATA_CONFIG,
+    get_metadata_path,
+    get_subdirectory_name,
+    resolve_subdirectory_path,
+)
+from .metadata_migration import detect_legacy_format, migrate_legacy_metadata, migrate_plate_metadata
+from .roi import (
+    ROI,
+    PolygonShape,
+    PolylineShape,
+    MaskShape,
+    PointShape,
+    EllipseShape,
+    extract_rois_from_labeled_mask,
+    load_rois_from_json,
+    load_rois_from_zip,
+    materialize_rois,
+)
+from .streaming import StreamingBackend
+from .streaming_constants import StreamingDataType, NapariShapeType
+
+__all__ = [
+    "Backend",
+    "MemoryType",
+    "TransportMode",
+    "FileFormat",
+    "DEFAULT_IMAGE_EXTENSIONS",
+    "BackendBase",
+    "DataSink",
+    "DataSource",
+    "ReadOnlyBackend",
+    "StorageBackend",
+    "StreamingBackend",
+    "storage_registry",
+    "reset_memory_backend",
+    "ensure_storage_registry",
+    "get_backend",
+    "get_backend_instance",
+    "cleanup_backend_connections",
+    "cleanup_all_backends",
+    "register_cleanup_callback",
+    "STORAGE_BACKENDS",
+    "DiskStorageBackend",
+    "MemoryStorageBackend",
+    "FileManager",
+    "file_lock",
+    "atomic_write_json",
+    "atomic_update_json",
+    "FileLockError",
+    "FileLockTimeoutError",
+    "AtomicMetadataWriter",
+    "MetadataWriteError",
+    "METADATA_CONFIG",
+    "get_metadata_path",
+    "get_subdirectory_name",
+    "resolve_subdirectory_path",
+    "detect_legacy_format",
+    "migrate_legacy_metadata",
+    "migrate_plate_metadata",
+    "ROI",
+    "PolygonShape",
+    "PolylineShape",
+    "MaskShape",
+    "PointShape",
+    "EllipseShape",
+    "extract_rois_from_labeled_mask",
+    "load_rois_from_json",
+    "load_rois_from_zip",
+    "materialize_rois",
+    "StreamingDataType",
+    "NapariShapeType",
+    "NapariStreamingBackend",
+    "FijiStreamingBackend",
+    "ZarrStorageBackend",
+    "OMEROLocalBackend",
+    "OMEROFileFormatRegistry",
+]
+
+_LAZY_BACKEND_REGISTRY = {
+    "NapariStreamingBackend": ("polystore.napari_stream", "NapariStreamingBackend"),
+    "FijiStreamingBackend": ("polystore.fiji_stream", "FijiStreamingBackend"),
+    "ZarrStorageBackend": ("polystore.zarr", "ZarrStorageBackend"),
+    "OMEROLocalBackend": ("polystore.omero_local", "OMEROLocalBackend"),
+    "OMEROFileFormatRegistry": ("polystore.omero_local", "OMEROFileFormatRegistry"),
+}
+
+
+def __getattr__(name):
+    """Lazy import of optional/extra backend classes."""
+    if name not in _LAZY_BACKEND_REGISTRY:
+        raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
+
+    if os.getenv("POLYSTORE_SUBPROCESS_NO_GPU") == "1":
+        class PlaceholderBackend:
+            pass
+        PlaceholderBackend.__name__ = name
+        PlaceholderBackend.__qualname__ = name
+        return PlaceholderBackend
+
+    module_path, class_name = _LAZY_BACKEND_REGISTRY[name]
+    import importlib
+    module = importlib.import_module(module_path)
+    return getattr(module, class_name)

polystore/async_init.py

@@ -0,0 +1,65 @@
+"""
+Async initialization for I/O backends.
+
+This module provides background loading of GPU-heavy dependencies
+to avoid blocking during import while ensuring they're ready when needed.
+"""
+import logging
+import threading
+
+logger = logging.getLogger(__name__)
+
+_init_complete = threading.Event()
+_init_thread = None
+
+
+def _background_init():
+    """
+    Background thread target for async initialization.
+
+    NOTE: Due to Python's GIL, importing heavy modules in a background thread
+    still blocks the main thread during imports. Therefore, we use lazy initialization
+    for GPU libraries and storage backends - they'll be loaded on first use.
+
+    Currently this is a no-op placeholder. GPU registry initialization happens
+    lazily on first use, not during startup.
+    """
+    try:
+        logger.info("Background I/O initialization complete (lazy mode - no-op)")
+    except Exception as e:
+        logger.error(f"Background I/O initialization failed: {e}")
+    finally:
+        _init_complete.set()
+
+
+def start_async_initialization():
+    """
+    Start background initialization of GPU-heavy dependencies.
+    
+    Call this during application startup (GUI/CLI) to pre-load
+    GPU libraries without blocking. Safe to call multiple times.
+    """
+    global _init_thread
+    
+    if _init_thread is None:
+        _init_thread = threading.Thread(
+            target=_background_init,
+            daemon=True,
+            name="io-async-init"
+        )
+        _init_thread.start()
+        logger.info("Started background I/O initialization")
+
+
+def wait_for_initialization(timeout: float = 60.0) -> bool:
+    """
+    Wait for background initialization to complete.
+    
+    Args:
+        timeout: Maximum seconds to wait
+        
+    Returns:
+        True if completed, False if timeout
+    """
+    return _init_complete.wait(timeout)
+

polystore/atomic.py

@@ -0,0 +1,194 @@
+"""
+Atomic file operations with locking.
+
+Provides utilities for atomic read-modify-write operations with file locking
+to prevent concurrency issues in multiprocessing environments.
+"""
+
+import json
+import logging
+import os
+import tempfile
+import time
+from contextlib import contextmanager
+from dataclasses import dataclass
+from pathlib import Path
+from typing import Any, Callable, Dict, Optional, TypeVar, Union
+
+# Cross-platform file locking
+try:
+    import fcntl
+    FCNTL_AVAILABLE = True
+except ImportError:
+    # Windows compatibility - use portalocker
+    import portalocker
+    FCNTL_AVAILABLE = False
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar('T')
+
+
+@dataclass(frozen=True)
+class LockConfig:
+    """Configuration constants for file locking operations."""
+    DEFAULT_TIMEOUT: float = 30.0
+    DEFAULT_POLL_INTERVAL: float = 0.1
+    LOCK_SUFFIX: str = '.lock'
+    TEMP_PREFIX: str = '.tmp'
+    JSON_INDENT: int = 2
+
+
+LOCK_CONFIG = LockConfig()
+
+
+class FileLockError(Exception):
+    """Raised when file locking operations fail."""
+    pass
+
+
+class FileLockTimeoutError(FileLockError):
+    """Raised when file lock acquisition times out."""
+    pass
+
+
+@contextmanager
+def file_lock(
+    lock_path: Union[str, Path],
+    timeout: float = LOCK_CONFIG.DEFAULT_TIMEOUT,
+    poll_interval: float = LOCK_CONFIG.DEFAULT_POLL_INTERVAL
+):
+    """Context manager for exclusive file locking."""
+    lock_path = Path(lock_path)
+    lock_path.parent.mkdir(parents=True, exist_ok=True)
+
+    lock_fd = None
+    try:
+        lock_fd = _acquire_lock_with_timeout(lock_path, timeout, poll_interval)
+        yield
+    except FileLockTimeoutError:
+        raise
+    except Exception as e:
+        raise FileLockError(f"File lock operation failed for {lock_path}: {e}") from e
+    finally:
+        _cleanup_lock(lock_fd, lock_path)
+
+
+def _acquire_lock_with_timeout(lock_path: Path, timeout: float, poll_interval: float) -> int:
+    """Acquire file lock with timeout and return file descriptor."""
+    deadline = time.time() + timeout
+
+    while time.time() < deadline:
+        if lock_fd := _try_acquire_lock(lock_path):
+            return lock_fd
+        time.sleep(poll_interval)
+
+    raise FileLockTimeoutError(f"Failed to acquire lock {lock_path} within {timeout}s")
+
+
+def _try_acquire_lock(lock_path: Path) -> Optional[int]:
+    """Try to acquire lock once, return fd or None."""
+    try:
+        lock_fd = os.open(str(lock_path), os.O_CREAT | os.O_WRONLY | os.O_TRUNC)
+        if FCNTL_AVAILABLE:
+            fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
+        else:
+            # Windows: use portalocker
+            portalocker.lock(lock_fd, portalocker.LOCK_EX | portalocker.LOCK_NB)
+        logger.debug(f"Acquired file lock: {lock_path}")
+        return lock_fd
+    except (OSError, IOError):
+        return None
+
+
+def _cleanup_lock(lock_fd: Optional[int], lock_path: Path) -> None:
+    """Clean up file lock resources."""
+    if lock_fd is not None:
+        try:
+            if FCNTL_AVAILABLE:
+                fcntl.flock(lock_fd, fcntl.LOCK_UN)
+            else:
+                # Windows: use portalocker
+                portalocker.unlock(lock_fd)
+            os.close(lock_fd)
+            logger.debug(f"Released file lock: {lock_path}")
+        except Exception as e:
+            logger.warning(f"Error releasing lock {lock_path}: {e}")
+
+    if lock_path.exists():
+        try:
+            lock_path.unlink()
+        except Exception as e:
+            logger.warning(f"Error removing lock file {lock_path}: {e}")
+
+
+def atomic_write_json(
+    file_path: Union[str, Path],
+    data: Dict[str, Any],
+    indent: int = LOCK_CONFIG.JSON_INDENT,
+    ensure_directory: bool = True
+) -> None:
+    """Atomically write JSON data to file using temporary file + rename."""
+    file_path = Path(file_path)
+
+    if ensure_directory:
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+
+    try:
+        tmp_path = _write_to_temp_file(file_path, data, indent)
+        # Use os.replace() instead of os.rename() for atomic replacement on all platforms
+        # os.rename() fails on Windows if destination exists, os.replace() works on both Unix and Windows
+        os.replace(tmp_path, str(file_path))
+        logger.debug(f"Atomically wrote JSON to {file_path}")
+    except Exception as e:
+        raise FileLockError(f"Atomic JSON write failed for {file_path}: {e}") from e
+
+
+def _write_to_temp_file(file_path: Path, data: Dict[str, Any], indent: int) -> str:
+    """Write data to temporary file and return path."""
+    with tempfile.NamedTemporaryFile(
+        mode='w',
+        dir=file_path.parent,
+        prefix=f"{LOCK_CONFIG.TEMP_PREFIX}{file_path.name}",
+        suffix='.json',
+        delete=False
+    ) as tmp_file:
+        json.dump(data, tmp_file, indent=indent)
+        tmp_file.flush()
+        os.fsync(tmp_file.fileno())
+        return tmp_file.name
+
+
+def atomic_update_json(
+    file_path: Union[str, Path],
+    update_func: Callable[[Optional[Dict[str, Any]]], Dict[str, Any]],
+    lock_timeout: float = LOCK_CONFIG.DEFAULT_TIMEOUT,
+    default_data: Optional[Dict[str, Any]] = None
+) -> None:
+    """Atomically update JSON file using read-modify-write with file locking."""
+    file_path = Path(file_path)
+    lock_path = file_path.with_suffix(f'{file_path.suffix}{LOCK_CONFIG.LOCK_SUFFIX}')
+
+    with file_lock(lock_path, timeout=lock_timeout):
+        current_data = _read_json_or_default(file_path, default_data)
+
+        try:
+            updated_data = update_func(current_data)
+        except Exception as e:
+            raise FileLockError(f"Update function failed for {file_path}: {e}") from e
+
+        atomic_write_json(file_path, updated_data)
+        logger.debug(f"Atomically updated JSON file: {file_path}")
+
+
+def _read_json_or_default(file_path: Path, default_data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
+    """Read JSON file or return default data if file doesn't exist or is invalid."""
+    if not file_path.exists():
+        return default_data
+
+    try:
+        with open(file_path, 'r') as f:
+            return json.load(f)
+    except (json.JSONDecodeError, IOError) as e:
+        logger.warning(f"Failed to read {file_path}, using default: {e}")
+        return default_data

polystore/backend_registry.py

@@ -0,0 +1,160 @@
+"""
+Storage backend metaclass registration system.
+
+Eliminates hardcoded backend registration by using metaclass auto-registration.
+Backends are automatically discovered and registered when their classes are defined.
+"""
+
+import logging
+from typing import Callable, Dict, List
+
+from .base import BackendBase, DataSink
+
+logger = logging.getLogger(__name__)
+
+_backend_instances: Dict[str, DataSink] = {}
+_cleanup_callbacks: List[Callable[[], None]] = []
+
+# Registry auto-created by AutoRegisterMeta on BackendBase
+# Includes both StorageBackend (read-write) and ReadOnlyBackend (read-only) subclasses
+STORAGE_BACKENDS = BackendBase.__registry__
+
+
+def get_backend_instance(backend_type: str) -> DataSink:
+    """
+    Get backend instance by type with lazy instantiation.
+
+    Args:
+        backend_type: Backend type identifier (e.g., 'disk', 'memory')
+
+    Returns:
+        Backend instance
+
+    Raises:
+        KeyError: If backend type not registered
+        RuntimeError: If backend instantiation fails
+    """
+    if hasattr(backend_type, "value"):
+        backend_type = backend_type.value
+    backend_type = str(backend_type).lower()
+
+    # Return cached instance if available
+    if backend_type in _backend_instances:
+        return _backend_instances[backend_type]
+
+    # Get backend class from registry
+    if backend_type not in STORAGE_BACKENDS:
+        raise KeyError(f"Backend type '{backend_type}' not registered. "
+                      f"Available backends: {list(STORAGE_BACKENDS.keys())}")
+
+    backend_class = STORAGE_BACKENDS[backend_type]
+
+    try:
+        # Create and cache instance
+        instance = backend_class()
+        _backend_instances[backend_type] = instance
+        logger.debug(f"Created instance for backend '{backend_type}'")
+        return instance
+    except Exception as e:
+        raise RuntimeError(f"Failed to instantiate backend '{backend_type}': {e}") from e
+
+
+def create_storage_registry() -> Dict[str, DataSink]:
+    """
+    Create storage registry with all registered backends.
+
+    Returns:
+        Dictionary mapping backend types to instances
+    """
+    # Trigger discovery of all backends in polystore package
+    # This imports all backend modules (disk, memory, zarr, napari_stream, fiji_stream, etc.)
+    # and registers them via metaclass
+    STORAGE_BACKENDS._discover()
+    logger.debug("Triggered backend discovery via LazyDiscoveryDict")
+
+    # Backends that require context-specific initialization (e.g., plate_root)
+    # These are registered lazily when needed, not at startup
+    SKIP_BACKENDS = {'virtual_workspace', 'omero_local'}
+
+    registry = {}
+    for backend_type in STORAGE_BACKENDS.keys():
+        # Skip backends that need context-specific initialization
+        if backend_type in SKIP_BACKENDS:
+            logger.debug(f"Skipping backend '{backend_type}' - requires context-specific initialization")
+            continue
+
+        try:
+            registry[backend_type] = get_backend_instance(backend_type)
+        except Exception as e:
+            logger.warning(f"Failed to create instance for backend '{backend_type}': {e}")
+            continue
+
+    logger.info(f"Created storage registry with {len(registry)} backends: {list(registry.keys())}")
+    return registry
+
+
+def register_cleanup_callback(callback: Callable[[], None]) -> None:
+    """Register an integration-specific cleanup callback."""
+    _cleanup_callbacks.append(callback)
+
+
+def cleanup_backend_connections() -> None:
+    """
+    Clean up backend connections without affecting persistent resources.
+
+    For napari streaming backend, this cleans up ZeroMQ connections but
+    leaves the napari window open for future use.
+    """
+    import os
+
+    # Check if we're running in test mode
+    is_test_mode = (
+        'pytest' in os.environ.get('_', '') or
+        'PYTEST_CURRENT_TEST' in os.environ or
+        any('pytest' in arg for arg in __import__('sys').argv)
+    )
+
+    for backend_type, instance in _backend_instances.items():
+        # Use targeted cleanup for napari streaming to preserve window
+        if hasattr(instance, 'cleanup_connections'):
+            try:
+                instance.cleanup_connections()
+                logger.debug(f"Cleaned up connections for backend '{backend_type}'")
+            except Exception as e:
+                logger.warning(f"Failed to cleanup connections for backend '{backend_type}': {e}")
+        elif hasattr(instance, 'cleanup') and backend_type != 'napari_stream':
+            try:
+                instance.cleanup()
+                logger.debug(f"Cleaned up backend '{backend_type}'")
+            except Exception as e:
+                logger.warning(f"Failed to cleanup backend '{backend_type}': {e}")
+
+    # In test mode, also stop viewer processes to allow pytest to exit
+    if is_test_mode:
+        for callback in list(_cleanup_callbacks):
+            try:
+                callback()
+            except Exception as e:
+                logger.warning(f"Cleanup callback failed: {e}")
+
+    logger.info(f"Backend connections cleaned up ({'test mode' if is_test_mode else 'viewer windows preserved'})")
+
+
+def cleanup_all_backends() -> None:
+    """
+    Clean up all cached backend instances completely.
+
+    This is for full shutdown - clears instance cache and calls full cleanup.
+    Use cleanup_backend_connections() for test cleanup to preserve napari window.
+    """
+    for backend_type, instance in _backend_instances.items():
+        if hasattr(instance, 'cleanup'):
+            try:
+                instance.cleanup()
+                logger.debug(f"Cleaned up backend '{backend_type}'")
+            except Exception as e:
+                logger.warning(f"Failed to cleanup backend '{backend_type}': {e}")
+
+    _backend_instances.clear()
+    logger.info("All backend instances cleaned up")
+