polystore 0.1.2__py3-none-any.whl

polystore/__init__.py

@@ -0,0 +1,80 @@
+"""
+Polystore - Framework-agnostic multi-backend storage abstraction.
+
+Provides pluggable storage backends with multi-framework I/O support for
+NumPy, PyTorch, JAX, TensorFlow, CuPy, and Zarr.
+"""
+
+__version__ = "0.1.2"
+
+# Core abstractions
+from .base import (
+    DataSink,
+    DataSource,
+    StorageBackend,
+    VirtualBackend,
+    ReadOnlyBackend,
+)
+
+# Concrete backends
+from .memory import MemoryBackend
+from .disk import DiskBackend
+
+# Optional backends
+try:
+    from .zarr import ZarrBackend
+except ImportError:
+    ZarrBackend = None
+
+# File manager
+from .filemanager import FileManager
+
+# Registry
+from .backend_registry import BackendRegistry, create_storage_registry
+
+# Atomic operations
+from .atomic import atomic_write, atomic_write_json
+
+# Exceptions
+from .exceptions import (
+    StorageError,
+    StorageResolutionError,
+    BackendNotFoundError,
+    UnsupportedFormatError,
+)
+
+# Streaming (optional)
+try:
+    from .streaming import StreamingBackend
+except ImportError:
+    StreamingBackend = None
+
+__all__ = [
+    # Version
+    "__version__",
+    # Core abstractions
+    "DataSink",
+    "DataSource",
+    "StorageBackend",
+    "VirtualBackend",
+    "ReadOnlyBackend",
+    # Backends
+    "MemoryBackend",
+    "DiskBackend",
+    "ZarrBackend",
+    # File manager
+    "FileManager",
+    # Registry
+    "BackendRegistry",
+    # Atomic operations
+    "atomic_write",
+    "atomic_write_json",
+    # Exceptions
+    "StorageError",
+    "StorageResolutionError",
+    "BackendNotFoundError",
+    "UnsupportedFormatError",
+    # Streaming
+    "StreamingBackend",
+]
+

polystore/atomic.py

@@ -0,0 +1,249 @@
+"""
+Atomic file operations with locking for OpenHCS.
+
+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}")
+
+
+@contextmanager
+def atomic_write(file_path: Union[str, Path], mode: str = 'w', ensure_directory: bool = True):
+    """
+    Context manager for atomic file writes.
+    
+    Writes to a temporary file first, then renames to the target path.
+    This ensures the operation is atomic - either fully succeeds or fails
+    without leaving partial writes.
+    
+    Args:
+        file_path: Target file path
+        mode: File mode ('w' for text, 'wb' for binary)
+        ensure_directory: Create parent directory if it doesn't exist
+        
+    Example:
+        with atomic_write("output.txt") as f:
+            f.write("data")
+    """
+    file_path = Path(file_path)
+    
+    if ensure_directory:
+        file_path.parent.mkdir(parents=True, exist_ok=True)
+    
+    # Create temporary file in same directory
+    with tempfile.NamedTemporaryFile(
+        mode=mode,
+        dir=file_path.parent,
+        prefix=f"{LOCK_CONFIG.TEMP_PREFIX}{file_path.name}",
+        delete=False
+    ) as tmp_file:
+        tmp_path = tmp_file.name
+        try:
+            yield tmp_file
+            tmp_file.flush()
+            os.fsync(tmp_file.fileno())
+        except Exception:
+            # Clean up temp file on error
+            try:
+                os.unlink(tmp_path)
+            except Exception:
+                pass
+            raise
+    
+    # Atomically replace target file
+    try:
+        os.replace(tmp_path, str(file_path))
+        logger.debug(f"Atomically wrote to {file_path}")
+    except Exception as e:
+        try:
+            os.unlink(tmp_path)
+        except Exception:
+            pass
+        raise FileLockError(f"Atomic write failed for {file_path}: {e}") from 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,179 @@
+"""
+Storage backend metaclass registration system.
+
+Backends are automatically discovered and registered when their classes are defined.
+"""
+
+import logging
+from typing import Dict
+from .base import BackendBase, DataSink
+
+logger = logging.getLogger(__name__)
+
+_backend_instances: Dict[str, DataSink] = {}
+
+
+def _get_storage_backends() -> Dict:
+    """Get the storage backends registry, ensuring it's initialized."""
+    # Import backends to trigger registration
+    from . import memory, disk
+    try:
+        from . import zarr
+    except ImportError:
+        pass  # Zarr not available
+    
+    # Registry auto-created by AutoRegisterMeta on BackendBase
+    return BackendBase.__registry__
+
+
+# Lazy access to registry
+STORAGE_BACKENDS = None
+
+
+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
+    """
+    backend_type = backend_type.lower()
+
+    # Return cached instance if available
+    if backend_type in _backend_instances:
+        return _backend_instances[backend_type]
+
+    # Get backend class from registry
+    storage_backends = _get_storage_backends()
+    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
+    """
+    # Get backends registry (triggers import and registration)
+    storage_backends = _get_storage_backends()
+
+    # Backends that require context-specific initialization (e.g., plate_root)
+    # These are registered lazily when needed, not at startup
+    SKIP_BACKENDS = {'virtual_workspace'}
+
+    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 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:
+        try:
+            from openhcs.runtime.napari_stream_visualizer import _cleanup_global_viewer
+            _cleanup_global_viewer()
+            logger.debug("Cleaned up napari viewer for test mode")
+        except ImportError:
+            pass  # napari not available
+        except Exception as e:
+            logger.warning(f"Failed to cleanup napari viewer: {e}")
+
+
+class BackendRegistry(dict):
+    """
+    Registry for storage backends.
+
+    This is a dictionary that automatically populates with available backends
+    when first accessed.
+    """
+
+    def __init__(self):
+        """Initialize the backend registry."""
+        super().__init__()
+        # Populate with available backends
+        self.update(create_storage_registry())
+
+
+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")
+
+
+