lamindb 1.1.1__py3-none-any.whl → 1.2.0__py3-none-any.whl

lamindb/core/storage/_tiledbsoma.py

@@ -12,8 +12,6 @@ from lamindb_setup.core._settings_storage import get_storage_region
 from lamindb_setup.core.upath import LocalPathClasses, create_path
 from packaging import version
 
-from lamindb.models import Artifact, Run
-
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
     from tiledbsoma import Collection as SOMACollection
@@ -21,12 +19,15 @@ if TYPE_CHECKING:
     from tiledbsoma import Measurement as SOMAMeasurement
     from upath import UPath
 
+    from lamindb.models.artifact import Artifact
+    from lamindb.models.run import Run
+
 
 def _load_h5ad_zarr(objpath: UPath):
-    from lamindb.core.loaders import load_anndata_zarr, load_h5ad
+    from lamindb.core.loaders import load_h5ad, load_zarr
 
     if objpath.is_dir():
-        adata = load_anndata_zarr(objpath)
+        adata = load_zarr(objpath, expected_type="anndata")
     else:
         # read only local in backed for now
         # in principle possible to read remote in backed also
@@ -134,9 +135,10 @@ def save_tiledbsoma_experiment(
     except ImportError as e:
         raise ImportError("Please install tiledbsoma: pip install tiledbsoma") from e
 
-    from lamindb.core._data import get_run
     from lamindb.core.storage.paths import auto_storage_key_from_artifact_uid
-    from lamindb.core.versioning import create_uid
+    from lamindb.models import Artifact
+    from lamindb.models._is_versioned import create_uid
+    from lamindb.models.artifact import get_run
 
     run = get_run(run)
 

lamindb/core/storage/_zarr.py

@@ -1,55 +1,134 @@
 from __future__ import annotations
 
 import warnings
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Literal
 
 import scipy.sparse as sparse
 import zarr
 from anndata import __version__ as anndata_version
 from anndata._io.specs import write_elem
-from anndata._io.specs.registry import get_spec
 from fsspec.implementations.local import LocalFileSystem
-from lamindb_setup.core.upath import create_mapper, infer_filesystem
+from lamin_utils import logger
+from lamindb_setup.core.upath import S3FSMap, create_mapper, infer_filesystem
 from packaging import version
 
+from lamindb.core._compat import with_package
+
 from ._anndata_sizes import _size_elem, _size_raw, size_adata
 
 if version.parse(anndata_version) < version.parse("0.11.0"):
-    from anndata._io import read_zarr
+    from anndata._io import read_zarr as read_anndata_zarr
 else:
-    from anndata.io import read_zarr
+    from anndata.io import read_zarr as read_anndata_zarr
 
 
 if TYPE_CHECKING:
     from anndata import AnnData
+    from fsspec import FSMap
     from lamindb_setup.core.types import UPathStr
 
+    from lamindb.core.types import ScverseDataStructures
 
-def zarr_is_adata(storepath: UPathStr) -> bool:
-    fs, storepath_str = infer_filesystem(storepath)
-    if isinstance(fs, LocalFileSystem):
-        # this is faster than through an fsspec mapper for local
-        open_obj = storepath_str
-    else:
-        open_obj = create_mapper(fs, storepath_str, check=True)
-    storage = zarr.open(open_obj, mode="r")
-    return get_spec(storage).encoding_type == "anndata"
 
-
-def load_anndata_zarr(storepath: UPathStr) -> AnnData:
+def create_zarr_open_obj(
+    storepath: UPathStr, *, check: bool = True
+) -> str | S3FSMap | FSMap:
+    """Creates the correct object that can be used to open a zarr file depending on local or remote location."""
     fs, storepath_str = infer_filesystem(storepath)
+
     if isinstance(fs, LocalFileSystem):
-        # this is faster than through an fsspec mapper for local
         open_obj = storepath_str
     else:
-        open_obj = create_mapper(fs, storepath_str, check=True)
-    adata = read_zarr(open_obj)
-    return adata
+        open_obj = create_mapper(fs, storepath_str, check=check)
+
+    return open_obj
+
+
+def _identify_zarr_type_from_storage(
+    storage: zarr.Group,
+) -> Literal["anndata", "mudata", "spatialdata", "unknown"]:
+    """Internal helper to identify zarr type from an open storage object."""
+    try:
+        if storage.attrs.get("encoding-type", "") == "anndata":
+            return "anndata"
+        elif storage.attrs.get("encoding-type", "") == "MuData":
+            return "mudata"
+        elif "spatialdata_attrs" in storage.attrs:
+            return "spatialdata"
+    except Exception as error:
+        logger.warning(f"an exception occurred {error}")
+    return "unknown"
+
+
+def identify_zarr_type(
+    storepath: UPathStr, *, check: bool = True
+) -> Literal["anndata", "mudata", "spatialdata", "unknown"]:
+    """Identify whether a zarr store is AnnData, SpatialData, or unknown type."""
+    # we can add these cheap suffix-based-checks later
+    # also need to check whether the .spatialdata.zarr suffix
+    # actually becomes a "standard"; currently we don't recognize it
+    # unlike ".anndata.zarr" in VALID_SUFFIXES
+    # suffixes = UPath(storepath).suffixes
+    # if ".spatialdata" in suffixes:
+    #     return "spatialdata"
+    # elif ".anndata" in suffixes:
+    #     return "anndata"
+
+    open_obj = create_zarr_open_obj(storepath, check=check)
+    try:
+        storage = zarr.open(open_obj, mode="r")
+        return _identify_zarr_type_from_storage(storage)
+    except Exception as error:
+        logger.warning(
+            f"an exception occured while trying to open the zarr store\n {error}"
+        )
+    return "unknown"
+
+
+def load_zarr(
+    storepath: UPathStr,
+    expected_type: Literal["anndata", "mudata", "spatialdata"] = None,
+) -> ScverseDataStructures:
+    """Loads a zarr store and returns the corresponding scverse data structure.
+
+    Args:
+        storepath: Path to the zarr store
+        expected_type: If provided, ensures the zarr store is of this type ("anndata", "mudata", "spatialdata")
+                       and raises ValueError if it's not
+    """
+    open_obj = create_zarr_open_obj(storepath, check=True)
+
+    # Open the storage once
+    try:
+        storage = zarr.open(open_obj, mode="r")
+    except Exception as error:
+        raise ValueError(f"Could not open zarr store: {error}") from None
+
+    actual_type = _identify_zarr_type_from_storage(storage)
+    if expected_type is not None and actual_type != expected_type:
+        raise ValueError(
+            f"Expected zarr store of type '{expected_type}', but found '{actual_type}'"
+        )
+
+    match actual_type:
+        case "anndata":
+            scverse_obj = read_anndata_zarr(open_obj)
+        case "mudata":
+            scverse_obj = with_package("mudata", lambda mod: mod.read_zarr(open_obj))
+        case "spatialdata":
+            scverse_obj = with_package(
+                "spatialdata", lambda mod: mod.read_zarr(open_obj)
+            )
+        case "unknown" | _:
+            raise ValueError(
+                "Unable to determine zarr store format and therefore cannot load Artifact."
+            )
+    return scverse_obj
 
 
 def write_adata_zarr(
     adata: AnnData, storepath: UPathStr, callback=None, chunks=None, **dataset_kwargs
-):
+) -> None:
     fs, storepath_str = infer_filesystem(storepath)
     store = create_mapper(fs, storepath_str, create=True)
 
@@ -65,7 +144,7 @@ def write_adata_zarr(
     adata_size = None
     cumulative_val = 0
 
-    def _cb(key_write: str | None = None):
+    def _report_progress(key_write: str | None = None):
         nonlocal adata_size
         nonlocal cumulative_val
 
@@ -91,9 +170,9 @@ def write_adata_zarr(
 
     def _write_elem_cb(f, k, elem, dataset_kwargs):
         write_elem(f, k, elem, dataset_kwargs=dataset_kwargs)
-        _cb(k)
+        _report_progress(k)
 
-    _cb(None)
+    _report_progress(None)
     with warnings.catch_warnings():
         warnings.filterwarnings("ignore", category=UserWarning, module="zarr")
 
@@ -114,4 +193,4 @@ def write_adata_zarr(
             )
         _write_elem_cb(f, "raw", adata.raw, dataset_kwargs=dataset_kwargs)
     # todo: fix size less than total at the end
-    _cb(None)
+    _report_progress(None)

lamindb/core/storage/objects.py

@@ -1,62 +1,97 @@
 from __future__ import annotations
 
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias
 
 from anndata import AnnData
 from pandas import DataFrame
 
+from lamindb.core._compat import (
+    with_package_obj,
+)
+from lamindb.core.types import ScverseDataStructures
+
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
-
-def _mudata_is_installed():
-    try:
-        import mudata  # noqa: F401c
-    except ImportError:
-        return False
-    return True
+SupportedDataTypes: TypeAlias = DataFrame | ScverseDataStructures
 
 
-def infer_suffix(dmem, adata_format: str | None = None):
+def infer_suffix(dmem: SupportedDataTypes, format: str | None = None):
     """Infer LaminDB storage file suffix from a data object."""
     if isinstance(dmem, AnnData):
-        if adata_format is not None:
-            if adata_format not in {"h5ad", "zarr", "anndata.zarr"}:
+        if format is not None:
+            if format not in {"h5ad", "zarr", "anndata.zarr"}:
                 raise ValueError(
                     "Error when specifying AnnData storage format, it should be"
-                    f" 'h5ad', 'zarr', not '{adata_format}'. Check 'format'"
+                    f" 'h5ad', 'zarr', not '{format}'. Check 'format'"
                     " or the suffix of 'key'."
                 )
-            return "." + adata_format
+            return "." + format
         return ".h5ad"
-    elif isinstance(dmem, DataFrame):
+
+    if isinstance(dmem, DataFrame):
         return ".parquet"
-    else:
-        if _mudata_is_installed():
-            from mudata import MuData
 
-            if isinstance(dmem, MuData):
-                return ".h5mu"
+    if with_package_obj(
+        dmem,
+        "MuData",
+        "mudata",
+        lambda obj: True,  # Just checking type, not calling any method
+    )[0]:
+        return ".h5mu"
+
+    has_spatialdata, spatialdata_suffix = with_package_obj(
+        dmem,
+        "SpatialData",
+        "spatialdata",
+        lambda obj: "."
+        + (
+            format
+            if format is not None and format in {"spatialdata.zarr", "zarr"}
+            else ".zarr"
+            if format is None
+            else (_ for _ in ()).throw(
+                ValueError(
+                    "Error when specifying SpatialData storage format, it should be"
+                    f" 'zarr', 'spatialdata.zarr', not '{format}'. Check 'format'"
+                    " or the suffix of 'key'."
+                )
+            )
+        ),
+    )
+    if has_spatialdata:
+        return spatialdata_suffix
+    else:
         raise NotImplementedError
 
 
-def write_to_disk(dmem, filepath: UPathStr):
+def write_to_disk(dmem: SupportedDataTypes, filepath: UPathStr) -> None:
+    """Writes the passed in memory data to disk to a specified path."""
     if isinstance(dmem, AnnData):
         suffix = PurePosixPath(filepath).suffix
         if suffix == ".h5ad":
             dmem.write_h5ad(filepath)
+            return
         elif suffix == ".zarr":
             dmem.write_zarr(filepath)
+            return
         else:
             raise NotImplementedError
-    elif isinstance(dmem, DataFrame):
+
+    if isinstance(dmem, DataFrame):
         dmem.to_parquet(filepath)
-    else:
-        if _mudata_is_installed():
-            from mudata import MuData
+        return
 
-            if isinstance(dmem, MuData):
-                dmem.write(filepath)
-                return
-        raise NotImplementedError
+    if with_package_obj(dmem, "MuData", "mudata", lambda obj: obj.write(filepath))[0]:
+        return
+
+    if with_package_obj(
+        dmem,
+        "SpatialData",
+        "spatialdata",
+        lambda obj: obj.write(filepath, overwrite=True),
+    )[0]:
+        return
+
+    raise NotImplementedError

lamindb/core/storage/paths.py

@@ -11,13 +11,14 @@ from lamindb_setup.core.upath import (
 )
 
 from lamindb.core._settings import settings
-from lamindb.models import Artifact, Storage
 
 if TYPE_CHECKING:
     from pathlib import Path
 
     from lamindb_setup.core.types import UPathStr
 
+    from lamindb.models.artifact import Artifact
+
 
 AUTO_KEY_PREFIX = ".lamindb/"
 
@@ -70,6 +71,8 @@ def attempt_accessing_path(
 ) -> tuple[UPath, StorageSettings]:
     # check whether the file is in the default db and whether storage
     # matches default storage
+    from lamindb.models import Storage
+
     if (
         artifact._state.db in ("default", None)
         and artifact.storage_id == settings._storage_settings.id

lamindb/core/types.py

@@ -1,3 +1,8 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, TypeVar
+
+from anndata import AnnData
 from lamindb_setup.core.types import UPathStr
 
 from lamindb.base.types import (
@@ -7,3 +12,8 @@ from lamindb.base.types import (
     StrField,
     TransformType,
 )
+
+MuData = TypeVar("MuData")
+SpatialData = TypeVar("SpatialData")
+
+ScverseDataStructures = AnnData | MuData | SpatialData