lamindb 1.1.0__py3-none-any.whl → 1.2a2__py3-none-any.whl

lamindb/core/storage/_backed_access.py

@@ -5,8 +5,6 @@ from typing import TYPE_CHECKING, Any, Callable
 
 from anndata._io.specs.registry import get_spec
 
-from lamindb.models import Artifact
-
 from ._anndata_accessor import AnnDataAccessor, StorageType, registry
 from ._pyarrow_dataset import _is_pyarrow_dataset, _open_pyarrow_dataset
 from ._tiledbsoma import _open_tiledbsoma
@@ -19,6 +17,8 @@ if TYPE_CHECKING:
     from tiledbsoma import Experiment as SOMAExperiment
     from upath import UPath
 
+    from lamindb.models.artifact import Artifact
+
 
 # this dynamically creates a subclass of a context manager class
 # and reassigns it to an instance of the superclass
@@ -70,9 +70,12 @@ def backed_access(
     artifact_or_filepath: Artifact | UPath,
     mode: str = "r",
     using_key: str | None = None,
+    **kwargs,
 ) -> (
     AnnDataAccessor | BackedAccessor | SOMACollection | SOMAExperiment | PyArrowDataset
 ):
+    from lamindb.models import Artifact
+
     if isinstance(artifact_or_filepath, Artifact):
         objectpath, _ = filepath_from_artifact(
             artifact_or_filepath, using_key=using_key
@@ -80,18 +83,22 @@ def backed_access(
     else:
         objectpath = artifact_or_filepath
     name = objectpath.name
-    suffix = objectpath.suffix
+    # ignore .gz, only check the real suffix
+    suffixes = objectpath.suffixes
+    suffix = (
+        suffixes[-2] if len(suffixes) > 1 and ".gz" in suffixes else objectpath.suffix
+    )
 
     if name == "soma" or suffix == ".tiledbsoma":
         if mode not in {"r", "w"}:
             raise ValueError("`mode` should be either 'r' or 'w' for tiledbsoma.")
-        return _open_tiledbsoma(objectpath, mode=mode)  # type: ignore
+        return _open_tiledbsoma(objectpath, mode=mode, **kwargs)  # type: ignore
     elif suffix in {".h5", ".hdf5", ".h5ad"}:
-        conn, storage = registry.open("h5py", objectpath, mode=mode)
+        conn, storage = registry.open("h5py", objectpath, mode=mode, **kwargs)
     elif suffix == ".zarr":
-        conn, storage = registry.open("zarr", objectpath, mode=mode)
+        conn, storage = registry.open("zarr", objectpath, mode=mode, **kwargs)
     elif _is_pyarrow_dataset(objectpath):
-        return _open_pyarrow_dataset(objectpath)
+        return _open_pyarrow_dataset(objectpath, **kwargs)
     else:
         raise ValueError(
             "The object should have .h5, .hdf5, .h5ad, .zarr, .tiledbsoma suffix "

lamindb/core/storage/_pyarrow_dataset.py

@@ -18,15 +18,30 @@ def _is_pyarrow_dataset(paths: UPath | list[UPath]) -> bool:
     # we don't check here that the filesystem is the same
     # but this is a requirement for pyarrow.dataset.dataset
     if isinstance(paths, list):
-        suffixes = {path.suffix for path in paths}
-    elif paths.is_file():
-        suffixes = {paths.suffix}
+        path_list = paths
+    elif paths.is_dir():
+        path_list = [path for path in paths.rglob("*") if path.suffix != ""]
     else:
-        suffixes = {path.suffix for path in paths.rglob("*") if path.suffix != ""}
-    return len(suffixes) == 1 and suffixes.pop() in PYARROW_SUFFIXES
-
-
-def _open_pyarrow_dataset(paths: UPath | list[UPath]) -> PyArrowDataset:
+        path_list = [paths]
+    suffix = None
+    for path in path_list:
+        path_suffixes = path.suffixes
+        # this doesn't work for externally gzipped files, REMOVE LATER
+        path_suffix = (
+            path_suffixes[-2]
+            if len(path_suffixes) > 1 and ".gz" in path_suffixes
+            else path.suffix
+        )
+        if path_suffix not in PYARROW_SUFFIXES:
+            return False
+        elif suffix is None:
+            suffix = path_suffix
+        elif path_suffix != suffix:
+            return False
+    return True
+
+
+def _open_pyarrow_dataset(paths: UPath | list[UPath], **kwargs) -> PyArrowDataset:
     if isinstance(paths, list):
         path0 = paths[0]
         if isinstance(path0, LocalPathClasses):
@@ -38,4 +53,4 @@ def _open_pyarrow_dataset(paths: UPath | list[UPath]) -> PyArrowDataset:
     else:
         paths_str, filesystem = paths.path, paths.fs
 
-    return pyarrow.dataset.dataset(paths_str, filesystem=filesystem)
+    return pyarrow.dataset.dataset(paths_str, filesystem=filesystem, **kwargs)

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,6 +19,9 @@ 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
@@ -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,14 +1,14 @@
 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 lamin_utils import logger
 from lamindb_setup.core.upath import create_mapper, infer_filesystem
 from packaging import version
 
@@ -25,15 +25,36 @@ if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
 
-def zarr_is_adata(storepath: UPathStr) -> bool:
+def identify_zarr_type(
+    storepath: UPathStr, *, check: bool = True
+) -> Literal["anndata", "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"
+
     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"
+        open_obj = create_mapper(fs, storepath_str, check=check)
+
+    try:
+        storage = zarr.open(open_obj, mode="r")
+        if "spatialdata_attrs" in storage.attrs:
+            return "spatialdata"
+        if storage.attrs.get("encoding-type", "") == "anndata":
+            return "anndata"
+    except Exception as error:
+        logger.warning(f"an exception occured {error}")
+    return "unknown"
 
 
 def load_anndata_zarr(storepath: UPathStr) -> AnnData:
@@ -65,7 +86,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 +112,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 +135,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,7 +1,7 @@
 from __future__ import annotations
 
 from pathlib import PurePosixPath
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, TypeAlias, TypeVar
 
 from anndata import AnnData
 from pandas import DataFrame
@@ -9,54 +9,87 @@ from pandas import DataFrame
 if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
+SpatialData = TypeVar("SpatialData")
+MuData = TypeVar("MuData")
 
-def _mudata_is_installed():
-    try:
-        import mudata  # noqa: F401c
-    except ImportError:
-        return False
-    return True
+SupportedDataTypes: TypeAlias = AnnData | DataFrame | MuData | SpatialData
 
 
-def infer_suffix(dmem, adata_format: str | None = None):
+def is_package_installed(package_name):
+    import importlib.util
+
+    spec = importlib.util.find_spec(package_name)
+    return spec is not 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 is_package_installed("mudata"):
+        from mudata import MuData
+
+        if isinstance(dmem, MuData):
+            return ".h5mu"
+
+    if is_package_installed("spatialdata"):
+        from spatialdata import SpatialData
+
+        if isinstance(dmem, SpatialData):
+            if format is not None:
+                if format not in {"spatialdata.zarr"}:
+                    raise ValueError(
+                        "Error when specifying SpatialData storage format, it should be"
+                        f" 'zarr', 'spatialdata.zarr', not '{format}'. Check 'format'"
+                        " or the suffix of 'key'."
+                    )
+                return "." + format
+            return ".zarr"
+    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 is_package_installed("mudata"):
+        from mudata import MuData
 
-            if isinstance(dmem, MuData):
-                dmem.write(filepath)
-                return
+        if isinstance(dmem, MuData):
+            dmem.write(filepath)
+            return
+
+    if is_package_installed("spatialdata"):
+        from spatialdata import SpatialData
+
+        if isinstance(dmem, SpatialData):
+            dmem.write(filepath, overwrite=True)
+            return
+    else:
         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/"
 
@@ -41,24 +42,21 @@ def auto_storage_key_from_artifact_uid(uid: str, suffix: str, is_dir: bool) -> s
     return storage_key
 
 
-def _safely_resolve(upath: UPath) -> UPath:
-    if upath.protocol in {"http", "https"}:
-        resolve_kwargs = {"follow_redirects": False}
-    else:
-        resolve_kwargs = {}
-    return upath.resolve(**resolve_kwargs)
-
-
 def check_path_is_child_of_root(path: UPathStr, root: UPathStr) -> bool:
     if fsspec.utils.get_protocol(str(path)) != fsspec.utils.get_protocol(str(root)):
         return False
-    path_upath = _safely_resolve(UPath(path))
-    root_upath = _safely_resolve(UPath(root))
+    path_upath = UPath(path)
+    root_upath = UPath(root)
     if path_upath.protocol == "s3":
         endpoint_path = path_upath.storage_options.get("endpoint_url", "")
         endpoint_root = root_upath.storage_options.get("endpoint_url", "")
         if endpoint_path != endpoint_root:
             return False
+    # we don't resolve http links because they can resolve into a different domain
+    # for example into a temporary url
+    if path_upath.protocol not in {"http", "https"}:
+        path_upath = path_upath.resolve()
+        root_upath = root_upath.resolve()
     # str is needed to eliminate UPath storage_options
     # which affect equality checks
     return UPath(str(root_upath)) in UPath(str(path_upath)).parents
@@ -73,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
@@ -134,7 +134,7 @@ def filepath_cache_key_from_artifact(
 
 
 def store_file_or_folder(
-    local_path: UPathStr, storage_path: UPath, print_progress: bool = True
+    local_path: UPathStr, storage_path: UPath, print_progress: bool = True, **kwargs
 ) -> None:
     """Store file or folder (localpath) at storagepath."""
     local_path = UPath(local_path)
@@ -155,7 +155,10 @@ def store_file_or_folder(
         else:
             create_folder = None
         storage_path.upload_from(
-            local_path, create_folder=create_folder, print_progress=print_progress
+            local_path,
+            create_folder=create_folder,
+            print_progress=print_progress,
+            **kwargs,
         )
     else:  # storage path is local
         if local_path.resolve().as_posix() == storage_path.resolve().as_posix():