Defer optional array runtimes at package import

Author: trissim

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "polystore"
-version = "0.1.28"
+version = "0.1.29"
 description = "Framework-agnostic multi-backend storage abstraction for ML and scientific computing"
 readme = "README.md"
 requires-python = ">=3.11"

src/polystore/__init__.py

@@ -3,6 +3,7 @@
 """
 
 from importlib.metadata import version as _distribution_version
+from typing import TYPE_CHECKING
 
 __version__ = _distribution_version("polystore")
 
@@ -35,7 +36,6 @@
     storage_registry,
 )
 from .constants import Backend, MemoryType
-from .disk import DiskBackend, DiskStorageBackend
 from .filemanager import FileManager
 from .formats import DEFAULT_IMAGE_EXTENSIONS, FileFormat
 from .memory import MemoryBackend, MemoryStorageBackend
@@ -68,6 +68,25 @@
 from .streaming_constants import NapariShapeType, StreamingDataType
 from .virtual_workspace import SourcePixelRef
 
+if TYPE_CHECKING:
+    from .disk import DiskBackend, DiskStorageBackend
+
+
+def __getattr__(name: str):
+    """Load disk implementations only when their public exports are requested."""
+
+    if name not in {"DiskBackend", "DiskStorageBackend"}:
+        raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+    from .disk import DiskBackend, DiskStorageBackend
+
+    globals().update(
+        DiskBackend=DiskBackend,
+        DiskStorageBackend=DiskStorageBackend,
+    )
+    return globals()[name]
+
+
 __all__ = [
     "Backend",
     "MemoryType",

tests/test_lazy_package_exports.py

@@ -0,0 +1,70 @@
+"""Package exports must not initialize optional array runtimes eagerly."""
+
+import json
+import subprocess
+import sys
+from pathlib import Path
+
+
+def _fresh_python(source: str) -> dict:
+    result = subprocess.run(
+        [sys.executable, "-c", source],
+        cwd=Path(__file__).resolve().parents[1],
+        check=True,
+        capture_output=True,
+        text=True,
+    )
+    return json.loads(result.stdout)
+
+
+def test_constants_import_does_not_load_disk_or_array_frameworks():
+    loaded = _fresh_python(
+        "\n".join(
+            (
+                "import json",
+                "import sys",
+                "from polystore.constants import Backend",
+                "assert Backend.DISK.value == 'disk'",
+                "names = ('polystore.disk', 'tensorflow', 'torch', 'cupy', 'jax')",
+                "print(json.dumps({name: name in sys.modules for name in names}))",
+            )
+        )
+    )
+
+    assert loaded == {
+        "polystore.disk": False,
+        "tensorflow": False,
+        "torch": False,
+        "cupy": False,
+        "jax": False,
+    }
+
+
+def test_root_disk_exports_load_and_cache_the_declared_classes():
+    loaded = _fresh_python(
+        "\n".join(
+            (
+                "import json",
+                "import polystore",
+                "disk_loaded_before = 'polystore.disk' in __import__('sys').modules",
+                "from polystore import DiskBackend, DiskStorageBackend",
+                "from polystore.disk import DiskBackend as DeclaredDiskBackend",
+                "from polystore.disk import DiskStorageBackend as DeclaredDiskStorageBackend",
+                "print(json.dumps({",
+                "    'disk_loaded_before': disk_loaded_before,",
+                "    'backend_identity': DiskBackend is DeclaredDiskBackend,",
+                "    'storage_identity': DiskStorageBackend is DeclaredDiskStorageBackend,",
+                "    'backend_cached': polystore.DiskBackend is DiskBackend,",
+                "    'storage_cached': polystore.DiskStorageBackend is DiskStorageBackend,",
+                "}))",
+            )
+        )
+    )
+
+    assert loaded == {
+        "disk_loaded_before": False,
+        "backend_identity": True,
+        "storage_identity": True,
+        "backend_cached": True,
+        "storage_cached": True,
+    }