polystore 0.1.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
polystore/__init__.py ADDED
@@ -0,0 +1,138 @@
1
+ """
2
+ Polystore package exports.
3
+ """
4
+
5
+ __version__ = "0.1.4"
6
+
7
+ import os
8
+
9
+ from .atomic import file_lock, atomic_write_json, atomic_update_json, FileLockError, FileLockTimeoutError
10
+ from .backend_registry import (
11
+ get_backend_instance,
12
+ cleanup_backend_connections,
13
+ cleanup_all_backends,
14
+ register_cleanup_callback,
15
+ STORAGE_BACKENDS,
16
+ )
17
+ from .base import (
18
+ BackendBase,
19
+ DataSink,
20
+ DataSource,
21
+ ReadOnlyBackend,
22
+ StorageBackend,
23
+ storage_registry,
24
+ reset_memory_backend,
25
+ ensure_storage_registry,
26
+ get_backend,
27
+ )
28
+ from .constants import Backend, MemoryType, TransportMode
29
+ from .disk import DiskStorageBackend
30
+ from .filemanager import FileManager
31
+ from .formats import FileFormat, DEFAULT_IMAGE_EXTENSIONS
32
+ from .memory import MemoryStorageBackend
33
+ from .metadata_writer import (
34
+ AtomicMetadataWriter,
35
+ MetadataWriteError,
36
+ METADATA_CONFIG,
37
+ get_metadata_path,
38
+ get_subdirectory_name,
39
+ resolve_subdirectory_path,
40
+ )
41
+ from .metadata_migration import detect_legacy_format, migrate_legacy_metadata, migrate_plate_metadata
42
+ from .roi import (
43
+ ROI,
44
+ PolygonShape,
45
+ PolylineShape,
46
+ MaskShape,
47
+ PointShape,
48
+ EllipseShape,
49
+ extract_rois_from_labeled_mask,
50
+ load_rois_from_json,
51
+ load_rois_from_zip,
52
+ materialize_rois,
53
+ )
54
+ from .streaming import StreamingBackend
55
+ from .streaming_constants import StreamingDataType, NapariShapeType
56
+
57
+ __all__ = [
58
+ "Backend",
59
+ "MemoryType",
60
+ "TransportMode",
61
+ "FileFormat",
62
+ "DEFAULT_IMAGE_EXTENSIONS",
63
+ "BackendBase",
64
+ "DataSink",
65
+ "DataSource",
66
+ "ReadOnlyBackend",
67
+ "StorageBackend",
68
+ "StreamingBackend",
69
+ "storage_registry",
70
+ "reset_memory_backend",
71
+ "ensure_storage_registry",
72
+ "get_backend",
73
+ "get_backend_instance",
74
+ "cleanup_backend_connections",
75
+ "cleanup_all_backends",
76
+ "register_cleanup_callback",
77
+ "STORAGE_BACKENDS",
78
+ "DiskStorageBackend",
79
+ "MemoryStorageBackend",
80
+ "FileManager",
81
+ "file_lock",
82
+ "atomic_write_json",
83
+ "atomic_update_json",
84
+ "FileLockError",
85
+ "FileLockTimeoutError",
86
+ "AtomicMetadataWriter",
87
+ "MetadataWriteError",
88
+ "METADATA_CONFIG",
89
+ "get_metadata_path",
90
+ "get_subdirectory_name",
91
+ "resolve_subdirectory_path",
92
+ "detect_legacy_format",
93
+ "migrate_legacy_metadata",
94
+ "migrate_plate_metadata",
95
+ "ROI",
96
+ "PolygonShape",
97
+ "PolylineShape",
98
+ "MaskShape",
99
+ "PointShape",
100
+ "EllipseShape",
101
+ "extract_rois_from_labeled_mask",
102
+ "load_rois_from_json",
103
+ "load_rois_from_zip",
104
+ "materialize_rois",
105
+ "StreamingDataType",
106
+ "NapariShapeType",
107
+ "NapariStreamingBackend",
108
+ "FijiStreamingBackend",
109
+ "ZarrStorageBackend",
110
+ "OMEROLocalBackend",
111
+ "OMEROFileFormatRegistry",
112
+ ]
113
+
114
+ _LAZY_BACKEND_REGISTRY = {
115
+ "NapariStreamingBackend": ("polystore.napari_stream", "NapariStreamingBackend"),
116
+ "FijiStreamingBackend": ("polystore.fiji_stream", "FijiStreamingBackend"),
117
+ "ZarrStorageBackend": ("polystore.zarr", "ZarrStorageBackend"),
118
+ "OMEROLocalBackend": ("polystore.omero_local", "OMEROLocalBackend"),
119
+ "OMEROFileFormatRegistry": ("polystore.omero_local", "OMEROFileFormatRegistry"),
120
+ }
121
+
122
+
123
+ def __getattr__(name):
124
+ """Lazy import of optional/extra backend classes."""
125
+ if name not in _LAZY_BACKEND_REGISTRY:
126
+ raise AttributeError(f"module '{__name__}' has no attribute '{name}'")
127
+
128
+ if os.getenv("POLYSTORE_SUBPROCESS_NO_GPU") == "1":
129
+ class PlaceholderBackend:
130
+ pass
131
+ PlaceholderBackend.__name__ = name
132
+ PlaceholderBackend.__qualname__ = name
133
+ return PlaceholderBackend
134
+
135
+ module_path, class_name = _LAZY_BACKEND_REGISTRY[name]
136
+ import importlib
137
+ module = importlib.import_module(module_path)
138
+ return getattr(module, class_name)
@@ -0,0 +1,65 @@
1
+ """
2
+ Async initialization for I/O backends.
3
+
4
+ This module provides background loading of GPU-heavy dependencies
5
+ to avoid blocking during import while ensuring they're ready when needed.
6
+ """
7
+ import logging
8
+ import threading
9
+
10
+ logger = logging.getLogger(__name__)
11
+
12
+ _init_complete = threading.Event()
13
+ _init_thread = None
14
+
15
+
16
+ def _background_init():
17
+ """
18
+ Background thread target for async initialization.
19
+
20
+ NOTE: Due to Python's GIL, importing heavy modules in a background thread
21
+ still blocks the main thread during imports. Therefore, we use lazy initialization
22
+ for GPU libraries and storage backends - they'll be loaded on first use.
23
+
24
+ Currently this is a no-op placeholder. GPU registry initialization happens
25
+ lazily on first use, not during startup.
26
+ """
27
+ try:
28
+ logger.info("Background I/O initialization complete (lazy mode - no-op)")
29
+ except Exception as e:
30
+ logger.error(f"Background I/O initialization failed: {e}")
31
+ finally:
32
+ _init_complete.set()
33
+
34
+
35
+ def start_async_initialization():
36
+ """
37
+ Start background initialization of GPU-heavy dependencies.
38
+
39
+ Call this during application startup (GUI/CLI) to pre-load
40
+ GPU libraries without blocking. Safe to call multiple times.
41
+ """
42
+ global _init_thread
43
+
44
+ if _init_thread is None:
45
+ _init_thread = threading.Thread(
46
+ target=_background_init,
47
+ daemon=True,
48
+ name="io-async-init"
49
+ )
50
+ _init_thread.start()
51
+ logger.info("Started background I/O initialization")
52
+
53
+
54
+ def wait_for_initialization(timeout: float = 60.0) -> bool:
55
+ """
56
+ Wait for background initialization to complete.
57
+
58
+ Args:
59
+ timeout: Maximum seconds to wait
60
+
61
+ Returns:
62
+ True if completed, False if timeout
63
+ """
64
+ return _init_complete.wait(timeout)
65
+
polystore/atomic.py ADDED
@@ -0,0 +1,194 @@
1
+ """
2
+ Atomic file operations with locking.
3
+
4
+ Provides utilities for atomic read-modify-write operations with file locking
5
+ to prevent concurrency issues in multiprocessing environments.
6
+ """
7
+
8
+ import json
9
+ import logging
10
+ import os
11
+ import tempfile
12
+ import time
13
+ from contextlib import contextmanager
14
+ from dataclasses import dataclass
15
+ from pathlib import Path
16
+ from typing import Any, Callable, Dict, Optional, TypeVar, Union
17
+
18
+ # Cross-platform file locking
19
+ try:
20
+ import fcntl
21
+ FCNTL_AVAILABLE = True
22
+ except ImportError:
23
+ # Windows compatibility - use portalocker
24
+ import portalocker
25
+ FCNTL_AVAILABLE = False
26
+
27
+ logger = logging.getLogger(__name__)
28
+
29
+ T = TypeVar('T')
30
+
31
+
32
+ @dataclass(frozen=True)
33
+ class LockConfig:
34
+ """Configuration constants for file locking operations."""
35
+ DEFAULT_TIMEOUT: float = 30.0
36
+ DEFAULT_POLL_INTERVAL: float = 0.1
37
+ LOCK_SUFFIX: str = '.lock'
38
+ TEMP_PREFIX: str = '.tmp'
39
+ JSON_INDENT: int = 2
40
+
41
+
42
+ LOCK_CONFIG = LockConfig()
43
+
44
+
45
+ class FileLockError(Exception):
46
+ """Raised when file locking operations fail."""
47
+ pass
48
+
49
+
50
+ class FileLockTimeoutError(FileLockError):
51
+ """Raised when file lock acquisition times out."""
52
+ pass
53
+
54
+
55
+ @contextmanager
56
+ def file_lock(
57
+ lock_path: Union[str, Path],
58
+ timeout: float = LOCK_CONFIG.DEFAULT_TIMEOUT,
59
+ poll_interval: float = LOCK_CONFIG.DEFAULT_POLL_INTERVAL
60
+ ):
61
+ """Context manager for exclusive file locking."""
62
+ lock_path = Path(lock_path)
63
+ lock_path.parent.mkdir(parents=True, exist_ok=True)
64
+
65
+ lock_fd = None
66
+ try:
67
+ lock_fd = _acquire_lock_with_timeout(lock_path, timeout, poll_interval)
68
+ yield
69
+ except FileLockTimeoutError:
70
+ raise
71
+ except Exception as e:
72
+ raise FileLockError(f"File lock operation failed for {lock_path}: {e}") from e
73
+ finally:
74
+ _cleanup_lock(lock_fd, lock_path)
75
+
76
+
77
+ def _acquire_lock_with_timeout(lock_path: Path, timeout: float, poll_interval: float) -> int:
78
+ """Acquire file lock with timeout and return file descriptor."""
79
+ deadline = time.time() + timeout
80
+
81
+ while time.time() < deadline:
82
+ if lock_fd := _try_acquire_lock(lock_path):
83
+ return lock_fd
84
+ time.sleep(poll_interval)
85
+
86
+ raise FileLockTimeoutError(f"Failed to acquire lock {lock_path} within {timeout}s")
87
+
88
+
89
+ def _try_acquire_lock(lock_path: Path) -> Optional[int]:
90
+ """Try to acquire lock once, return fd or None."""
91
+ try:
92
+ lock_fd = os.open(str(lock_path), os.O_CREAT | os.O_WRONLY | os.O_TRUNC)
93
+ if FCNTL_AVAILABLE:
94
+ fcntl.flock(lock_fd, fcntl.LOCK_EX | fcntl.LOCK_NB)
95
+ else:
96
+ # Windows: use portalocker
97
+ portalocker.lock(lock_fd, portalocker.LOCK_EX | portalocker.LOCK_NB)
98
+ logger.debug(f"Acquired file lock: {lock_path}")
99
+ return lock_fd
100
+ except (OSError, IOError):
101
+ return None
102
+
103
+
104
+ def _cleanup_lock(lock_fd: Optional[int], lock_path: Path) -> None:
105
+ """Clean up file lock resources."""
106
+ if lock_fd is not None:
107
+ try:
108
+ if FCNTL_AVAILABLE:
109
+ fcntl.flock(lock_fd, fcntl.LOCK_UN)
110
+ else:
111
+ # Windows: use portalocker
112
+ portalocker.unlock(lock_fd)
113
+ os.close(lock_fd)
114
+ logger.debug(f"Released file lock: {lock_path}")
115
+ except Exception as e:
116
+ logger.warning(f"Error releasing lock {lock_path}: {e}")
117
+
118
+ if lock_path.exists():
119
+ try:
120
+ lock_path.unlink()
121
+ except Exception as e:
122
+ logger.warning(f"Error removing lock file {lock_path}: {e}")
123
+
124
+
125
+ def atomic_write_json(
126
+ file_path: Union[str, Path],
127
+ data: Dict[str, Any],
128
+ indent: int = LOCK_CONFIG.JSON_INDENT,
129
+ ensure_directory: bool = True
130
+ ) -> None:
131
+ """Atomically write JSON data to file using temporary file + rename."""
132
+ file_path = Path(file_path)
133
+
134
+ if ensure_directory:
135
+ file_path.parent.mkdir(parents=True, exist_ok=True)
136
+
137
+ try:
138
+ tmp_path = _write_to_temp_file(file_path, data, indent)
139
+ # Use os.replace() instead of os.rename() for atomic replacement on all platforms
140
+ # os.rename() fails on Windows if destination exists, os.replace() works on both Unix and Windows
141
+ os.replace(tmp_path, str(file_path))
142
+ logger.debug(f"Atomically wrote JSON to {file_path}")
143
+ except Exception as e:
144
+ raise FileLockError(f"Atomic JSON write failed for {file_path}: {e}") from e
145
+
146
+
147
+ def _write_to_temp_file(file_path: Path, data: Dict[str, Any], indent: int) -> str:
148
+ """Write data to temporary file and return path."""
149
+ with tempfile.NamedTemporaryFile(
150
+ mode='w',
151
+ dir=file_path.parent,
152
+ prefix=f"{LOCK_CONFIG.TEMP_PREFIX}{file_path.name}",
153
+ suffix='.json',
154
+ delete=False
155
+ ) as tmp_file:
156
+ json.dump(data, tmp_file, indent=indent)
157
+ tmp_file.flush()
158
+ os.fsync(tmp_file.fileno())
159
+ return tmp_file.name
160
+
161
+
162
+ def atomic_update_json(
163
+ file_path: Union[str, Path],
164
+ update_func: Callable[[Optional[Dict[str, Any]]], Dict[str, Any]],
165
+ lock_timeout: float = LOCK_CONFIG.DEFAULT_TIMEOUT,
166
+ default_data: Optional[Dict[str, Any]] = None
167
+ ) -> None:
168
+ """Atomically update JSON file using read-modify-write with file locking."""
169
+ file_path = Path(file_path)
170
+ lock_path = file_path.with_suffix(f'{file_path.suffix}{LOCK_CONFIG.LOCK_SUFFIX}')
171
+
172
+ with file_lock(lock_path, timeout=lock_timeout):
173
+ current_data = _read_json_or_default(file_path, default_data)
174
+
175
+ try:
176
+ updated_data = update_func(current_data)
177
+ except Exception as e:
178
+ raise FileLockError(f"Update function failed for {file_path}: {e}") from e
179
+
180
+ atomic_write_json(file_path, updated_data)
181
+ logger.debug(f"Atomically updated JSON file: {file_path}")
182
+
183
+
184
+ def _read_json_or_default(file_path: Path, default_data: Optional[Dict[str, Any]]) -> Optional[Dict[str, Any]]:
185
+ """Read JSON file or return default data if file doesn't exist or is invalid."""
186
+ if not file_path.exists():
187
+ return default_data
188
+
189
+ try:
190
+ with open(file_path, 'r') as f:
191
+ return json.load(f)
192
+ except (json.JSONDecodeError, IOError) as e:
193
+ logger.warning(f"Failed to read {file_path}, using default: {e}")
194
+ return default_data
@@ -0,0 +1,160 @@
1
+ """
2
+ Storage backend metaclass registration system.
3
+
4
+ Eliminates hardcoded backend registration by using metaclass auto-registration.
5
+ Backends are automatically discovered and registered when their classes are defined.
6
+ """
7
+
8
+ import logging
9
+ from typing import Callable, Dict, List
10
+
11
+ from .base import BackendBase, DataSink
12
+
13
+ logger = logging.getLogger(__name__)
14
+
15
+ _backend_instances: Dict[str, DataSink] = {}
16
+ _cleanup_callbacks: List[Callable[[], None]] = []
17
+
18
+ # Registry auto-created by AutoRegisterMeta on BackendBase
19
+ # Includes both StorageBackend (read-write) and ReadOnlyBackend (read-only) subclasses
20
+ STORAGE_BACKENDS = BackendBase.__registry__
21
+
22
+
23
+ def get_backend_instance(backend_type: str) -> DataSink:
24
+ """
25
+ Get backend instance by type with lazy instantiation.
26
+
27
+ Args:
28
+ backend_type: Backend type identifier (e.g., 'disk', 'memory')
29
+
30
+ Returns:
31
+ Backend instance
32
+
33
+ Raises:
34
+ KeyError: If backend type not registered
35
+ RuntimeError: If backend instantiation fails
36
+ """
37
+ if hasattr(backend_type, "value"):
38
+ backend_type = backend_type.value
39
+ backend_type = str(backend_type).lower()
40
+
41
+ # Return cached instance if available
42
+ if backend_type in _backend_instances:
43
+ return _backend_instances[backend_type]
44
+
45
+ # Get backend class from registry
46
+ if backend_type not in STORAGE_BACKENDS:
47
+ raise KeyError(f"Backend type '{backend_type}' not registered. "
48
+ f"Available backends: {list(STORAGE_BACKENDS.keys())}")
49
+
50
+ backend_class = STORAGE_BACKENDS[backend_type]
51
+
52
+ try:
53
+ # Create and cache instance
54
+ instance = backend_class()
55
+ _backend_instances[backend_type] = instance
56
+ logger.debug(f"Created instance for backend '{backend_type}'")
57
+ return instance
58
+ except Exception as e:
59
+ raise RuntimeError(f"Failed to instantiate backend '{backend_type}': {e}") from e
60
+
61
+
62
+ def create_storage_registry() -> Dict[str, DataSink]:
63
+ """
64
+ Create storage registry with all registered backends.
65
+
66
+ Returns:
67
+ Dictionary mapping backend types to instances
68
+ """
69
+ # Trigger discovery of all backends in polystore package
70
+ # This imports all backend modules (disk, memory, zarr, napari_stream, fiji_stream, etc.)
71
+ # and registers them via metaclass
72
+ STORAGE_BACKENDS._discover()
73
+ logger.debug("Triggered backend discovery via LazyDiscoveryDict")
74
+
75
+ # Backends that require context-specific initialization (e.g., plate_root)
76
+ # These are registered lazily when needed, not at startup
77
+ SKIP_BACKENDS = {'virtual_workspace', 'omero_local'}
78
+
79
+ registry = {}
80
+ for backend_type in STORAGE_BACKENDS.keys():
81
+ # Skip backends that need context-specific initialization
82
+ if backend_type in SKIP_BACKENDS:
83
+ logger.debug(f"Skipping backend '{backend_type}' - requires context-specific initialization")
84
+ continue
85
+
86
+ try:
87
+ registry[backend_type] = get_backend_instance(backend_type)
88
+ except Exception as e:
89
+ logger.warning(f"Failed to create instance for backend '{backend_type}': {e}")
90
+ continue
91
+
92
+ logger.info(f"Created storage registry with {len(registry)} backends: {list(registry.keys())}")
93
+ return registry
94
+
95
+
96
+ def register_cleanup_callback(callback: Callable[[], None]) -> None:
97
+ """Register an integration-specific cleanup callback."""
98
+ _cleanup_callbacks.append(callback)
99
+
100
+
101
+ def cleanup_backend_connections() -> None:
102
+ """
103
+ Clean up backend connections without affecting persistent resources.
104
+
105
+ For napari streaming backend, this cleans up ZeroMQ connections but
106
+ leaves the napari window open for future use.
107
+ """
108
+ import os
109
+
110
+ # Check if we're running in test mode
111
+ is_test_mode = (
112
+ 'pytest' in os.environ.get('_', '') or
113
+ 'PYTEST_CURRENT_TEST' in os.environ or
114
+ any('pytest' in arg for arg in __import__('sys').argv)
115
+ )
116
+
117
+ for backend_type, instance in _backend_instances.items():
118
+ # Use targeted cleanup for napari streaming to preserve window
119
+ if hasattr(instance, 'cleanup_connections'):
120
+ try:
121
+ instance.cleanup_connections()
122
+ logger.debug(f"Cleaned up connections for backend '{backend_type}'")
123
+ except Exception as e:
124
+ logger.warning(f"Failed to cleanup connections for backend '{backend_type}': {e}")
125
+ elif hasattr(instance, 'cleanup') and backend_type != 'napari_stream':
126
+ try:
127
+ instance.cleanup()
128
+ logger.debug(f"Cleaned up backend '{backend_type}'")
129
+ except Exception as e:
130
+ logger.warning(f"Failed to cleanup backend '{backend_type}': {e}")
131
+
132
+ # In test mode, also stop viewer processes to allow pytest to exit
133
+ if is_test_mode:
134
+ for callback in list(_cleanup_callbacks):
135
+ try:
136
+ callback()
137
+ except Exception as e:
138
+ logger.warning(f"Cleanup callback failed: {e}")
139
+
140
+ logger.info(f"Backend connections cleaned up ({'test mode' if is_test_mode else 'viewer windows preserved'})")
141
+
142
+
143
+ def cleanup_all_backends() -> None:
144
+ """
145
+ Clean up all cached backend instances completely.
146
+
147
+ This is for full shutdown - clears instance cache and calls full cleanup.
148
+ Use cleanup_backend_connections() for test cleanup to preserve napari window.
149
+ """
150
+ for backend_type, instance in _backend_instances.items():
151
+ if hasattr(instance, 'cleanup'):
152
+ try:
153
+ instance.cleanup()
154
+ logger.debug(f"Cleaned up backend '{backend_type}'")
155
+ except Exception as e:
156
+ logger.warning(f"Failed to cleanup backend '{backend_type}': {e}")
157
+
158
+ _backend_instances.clear()
159
+ logger.info("All backend instances cleaned up")
160
+