lamindb 0.76.7__py3-none-any.whl → 0.76.8__py3-none-any.whl

lamindb/core/storage/_tiledbsoma.py

@@ -1,204 +1,204 @@
-from __future__ import annotations
-
-from typing import TYPE_CHECKING, Literal
-
-from anndata import AnnData, read_h5ad
-from lamindb_setup import settings as setup_settings
-from lamindb_setup.core._settings_storage import get_storage_region
-from lamindb_setup.core.upath import LocalPathClasses, create_path
-from lnschema_core import Artifact, Run
-
-if TYPE_CHECKING:
-    from lamindb_setup.core.types import UPathStr
-    from tiledbsoma import Collection as SOMACollection
-    from tiledbsoma import Experiment as SOMAExperiment
-    from upath import UPath
-
-
-def _load_h5ad_zarr(objpath: UPath):
-    from lamindb.core.loaders import load_anndata_zarr, load_h5ad
-
-    if objpath.is_dir():
-        adata = load_anndata_zarr(objpath)
-    else:
-        # read only local in backed for now
-        # in principle possible to read remote in backed also
-        if isinstance(objpath, LocalPathClasses):
-            adata = read_h5ad(objpath.as_posix(), backed="r")
-        else:
-            adata = load_h5ad(objpath)
-    return adata
-
-
-def _tiledb_config_s3(storepath: UPath) -> dict:
-    region = get_storage_region(storepath)
-    tiledb_config = {"vfs.s3.region": region}
-    storage_options = storepath.storage_options
-    if "key" in storage_options:
-        tiledb_config["vfs.s3.aws_access_key_id"] = storage_options["key"]
-    if "secret" in storage_options:
-        tiledb_config["vfs.s3.aws_secret_access_key"] = storage_options["secret"]
-    if "token" in storage_options:
-        tiledb_config["vfs.s3.aws_session_token"] = storage_options["token"]
-
-    return tiledb_config
-
-
-def _open_tiledbsoma(
-    storepath: UPath, mode: Literal["r", "w"] = "r"
-) -> SOMACollection | SOMAExperiment:
-    try:
-        import tiledbsoma as soma
-    except ImportError as e:
-        raise ImportError("Please install tiledbsoma: pip install tiledbsoma") from e
-
-    storepath_str = storepath.as_posix()
-    if storepath.protocol == "s3":
-        ctx = soma.SOMATileDBContext(tiledb_config=_tiledb_config_s3(storepath))
-        # this is a strange bug
-        # for some reason iterdir futher gives incorrect results
-        # if cache is not invalidated
-        # instead of obs and ms it gives ms and ms in the list of names
-        storepath.fs.invalidate_cache()
-    else:
-        ctx = None
-
-    soma_objects = [obj.name for obj in storepath.iterdir()]
-    if "obs" in soma_objects and "ms" in soma_objects:
-        SOMAType = soma.Experiment
-    else:
-        SOMAType = soma.Collection
-    return SOMAType.open(storepath_str, mode=mode, context=ctx)
-
-
-def save_tiledbsoma_experiment(
-    # Artifact args
-    adatas: list[AnnData | UPathStr],
-    key: str | None = None,
-    description: str | None = None,
-    run: Run | None = None,
-    revises: Artifact | None = None,
-    # tiledbsoma.io.from_anndata args
-    measurement_name: str = "RNA",
-    obs_id_name: str = "obs_id",
-    var_id_name: str = "var_id",
-    append_obsm_varm: bool = False,
-    # additional keyword args for tiledbsoma.io.from_anndata
-    **kwargs,
-) -> Artifact:
-    """Write `AnnData` to `tiledbsoma.Experiment`.
-
-    Reads `AnnData` objects, writes them to `tiledbsoma.Experiment`, creates & saves an {class}`~lamindb.Artifact`.
-
-    Populates a column `lamin_run_uid` column in `obs` with the current `run.uid`.
-
-    Is based on `tiledbsoma.io.from_anndata
-    <https://tiledbsoma.readthedocs.io/en/latest/_autosummary/tiledbsoma.io.from_anndata.html>`__.
-
-    Args:
-        adatas: `AnnData` objects to write, in-memory or on-disk.
-        key: An optional key to reference the artifact.
-        description: A description.
-        run: The run that creates the artifact.
-        revises: `lamindb.Artifact` with `tiledbsoma.Experiment` to append to.
-        measurement_name: The name of the measurement to store data in `tiledbsoma.Experiment`.
-        obs_id_name: Which `AnnData` `obs` column to use for append mode.
-        var_id_name: Which `AnnData` `var` column to use for append mode.
-        append_obsm_varm: Whether to append `obsm` and `varm` in append mode .
-        **kwargs: Keyword arguments passed to `tiledbsoma.io.from_anndata`.
-    """
-    try:
-        import tiledbsoma as soma
-        import tiledbsoma.io as soma_io
-    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
-
-    run = get_run(run)
-
-    appending = revises is not None
-    if appending:
-        storepath = revises.path
-    else:
-        uid, _ = create_uid(n_full_id=20)
-        storage_key = auto_storage_key_from_artifact_uid(
-            uid, ".tiledbsoma", is_dir=True
-        )
-        storepath = setup_settings.storage.root / storage_key
-
-    if storepath.protocol == "s3":
-        ctx = soma.SOMATileDBContext(tiledb_config=_tiledb_config_s3(storepath))
-    else:
-        ctx = None
-
-    storepath = storepath.as_posix()
-
-    add_run_uid = True
-    if appending:
-        with soma.Experiment.open(storepath, mode="r", context=ctx) as store:
-            add_run_uid = "lamin_run_uid" in store["obs"].schema.names
-
-    if add_run_uid and run is None:
-        raise ValueError("Pass `run`")
-
-    adata_objects = []
-    for adata in adatas:
-        if isinstance(adata, AnnData):
-            if add_run_uid:
-                if adata.is_view:
-                    raise ValueError(
-                        "Can not write an `AnnData` view, please do `adata.copy()` before passing."
-                    )
-                else:
-                    adata.obs["lamin_run_uid"] = run.uid
-        else:
-            adata = _load_h5ad_zarr(create_path(adata))
-            if add_run_uid:
-                adata.obs["lamin_run_uid"] = run.uid
-        adata_objects.append(adata)
-
-    registration_mapping = kwargs.get("registration_mapping", None)
-    if registration_mapping is None and (appending or len(adata_objects) > 1):
-        registration_mapping = soma_io.register_anndatas(
-            experiment_uri=storepath if appending else None,
-            adatas=adata_objects,
-            measurement_name=measurement_name,
-            obs_field_name=obs_id_name,
-            var_field_name=var_id_name,
-            append_obsm_varm=append_obsm_varm,
-            context=ctx,
-        )
-
-    if registration_mapping is not None:
-        n_observations = len(registration_mapping.obs_axis.data)
-    else:  # happens only if not appending and only one adata passed
-        assert len(adata_objects) == 1  # noqa: S101
-        n_observations = adata_objects[0].n_obs
-
-    for adata_obj in adata_objects:
-        soma_io.from_anndata(
-            storepath,
-            adata_obj,
-            measurement_name,
-            context=ctx,
-            obs_id_name=obs_id_name,
-            var_id_name=var_id_name,
-            registration_mapping=registration_mapping,
-            **kwargs,
-        )
-
-    artifact = Artifact(
-        storepath,
-        key=key,
-        description=description,
-        run=run,
-        revises=revises,
-        _is_internal_call=True,
-    )
-    artifact.n_observations = n_observations
-    artifact._accessor = "tiledbsoma"
-
-    return artifact.save()
+from __future__ import annotations
+
+from typing import TYPE_CHECKING, Literal
+
+from anndata import AnnData, read_h5ad
+from lamindb_setup import settings as setup_settings
+from lamindb_setup.core._settings_storage import get_storage_region
+from lamindb_setup.core.upath import LocalPathClasses, create_path
+from lnschema_core import Artifact, Run
+
+if TYPE_CHECKING:
+    from lamindb_setup.core.types import UPathStr
+    from tiledbsoma import Collection as SOMACollection
+    from tiledbsoma import Experiment as SOMAExperiment
+    from upath import UPath
+
+
+def _load_h5ad_zarr(objpath: UPath):
+    from lamindb.core.loaders import load_anndata_zarr, load_h5ad
+
+    if objpath.is_dir():
+        adata = load_anndata_zarr(objpath)
+    else:
+        # read only local in backed for now
+        # in principle possible to read remote in backed also
+        if isinstance(objpath, LocalPathClasses):
+            adata = read_h5ad(objpath.as_posix(), backed="r")
+        else:
+            adata = load_h5ad(objpath)
+    return adata
+
+
+def _tiledb_config_s3(storepath: UPath) -> dict:
+    region = get_storage_region(storepath)
+    tiledb_config = {"vfs.s3.region": region}
+    storage_options = storepath.storage_options
+    if "key" in storage_options:
+        tiledb_config["vfs.s3.aws_access_key_id"] = storage_options["key"]
+    if "secret" in storage_options:
+        tiledb_config["vfs.s3.aws_secret_access_key"] = storage_options["secret"]
+    if "token" in storage_options:
+        tiledb_config["vfs.s3.aws_session_token"] = storage_options["token"]
+
+    return tiledb_config
+
+
+def _open_tiledbsoma(
+    storepath: UPath, mode: Literal["r", "w"] = "r"
+) -> SOMACollection | SOMAExperiment:
+    try:
+        import tiledbsoma as soma
+    except ImportError as e:
+        raise ImportError("Please install tiledbsoma: pip install tiledbsoma") from e
+
+    storepath_str = storepath.as_posix()
+    if storepath.protocol == "s3":
+        ctx = soma.SOMATileDBContext(tiledb_config=_tiledb_config_s3(storepath))
+        # this is a strange bug
+        # for some reason iterdir futher gives incorrect results
+        # if cache is not invalidated
+        # instead of obs and ms it gives ms and ms in the list of names
+        storepath.fs.invalidate_cache()
+    else:
+        ctx = None
+
+    soma_objects = [obj.name for obj in storepath.iterdir()]
+    if "obs" in soma_objects and "ms" in soma_objects:
+        SOMAType = soma.Experiment
+    else:
+        SOMAType = soma.Collection
+    return SOMAType.open(storepath_str, mode=mode, context=ctx)
+
+
+def save_tiledbsoma_experiment(
+    # Artifact args
+    adatas: list[AnnData | UPathStr],
+    key: str | None = None,
+    description: str | None = None,
+    run: Run | None = None,
+    revises: Artifact | None = None,
+    # tiledbsoma.io.from_anndata args
+    measurement_name: str = "RNA",
+    obs_id_name: str = "obs_id",
+    var_id_name: str = "var_id",
+    append_obsm_varm: bool = False,
+    # additional keyword args for tiledbsoma.io.from_anndata
+    **kwargs,
+) -> Artifact:
+    """Write `AnnData` to `tiledbsoma.Experiment`.
+
+    Reads `AnnData` objects, writes them to `tiledbsoma.Experiment`, creates & saves an {class}`~lamindb.Artifact`.
+
+    Populates a column `lamin_run_uid` column in `obs` with the current `run.uid`.
+
+    Is based on `tiledbsoma.io.from_anndata
+    <https://tiledbsoma.readthedocs.io/en/latest/_autosummary/tiledbsoma.io.from_anndata.html>`__.
+
+    Args:
+        adatas: `AnnData` objects to write, in-memory or on-disk.
+        key: An optional key to reference the artifact.
+        description: A description.
+        run: The run that creates the artifact.
+        revises: `lamindb.Artifact` with `tiledbsoma.Experiment` to append to.
+        measurement_name: The name of the measurement to store data in `tiledbsoma.Experiment`.
+        obs_id_name: Which `AnnData` `obs` column to use for append mode.
+        var_id_name: Which `AnnData` `var` column to use for append mode.
+        append_obsm_varm: Whether to append `obsm` and `varm` in append mode .
+        **kwargs: Keyword arguments passed to `tiledbsoma.io.from_anndata`.
+    """
+    try:
+        import tiledbsoma as soma
+        import tiledbsoma.io as soma_io
+    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
+
+    run = get_run(run)
+
+    appending = revises is not None
+    if appending:
+        storepath = revises.path
+    else:
+        uid, _ = create_uid(n_full_id=20)
+        storage_key = auto_storage_key_from_artifact_uid(
+            uid, ".tiledbsoma", is_dir=True
+        )
+        storepath = setup_settings.storage.root / storage_key
+
+    if storepath.protocol == "s3":
+        ctx = soma.SOMATileDBContext(tiledb_config=_tiledb_config_s3(storepath))
+    else:
+        ctx = None
+
+    storepath = storepath.as_posix()
+
+    add_run_uid = True
+    if appending:
+        with soma.Experiment.open(storepath, mode="r", context=ctx) as store:
+            add_run_uid = "lamin_run_uid" in store["obs"].schema.names
+
+    if add_run_uid and run is None:
+        raise ValueError("Pass `run`")
+
+    adata_objects = []
+    for adata in adatas:
+        if isinstance(adata, AnnData):
+            if add_run_uid:
+                if adata.is_view:
+                    raise ValueError(
+                        "Can not write an `AnnData` view, please do `adata.copy()` before passing."
+                    )
+                else:
+                    adata.obs["lamin_run_uid"] = run.uid
+        else:
+            adata = _load_h5ad_zarr(create_path(adata))
+            if add_run_uid:
+                adata.obs["lamin_run_uid"] = run.uid
+        adata_objects.append(adata)
+
+    registration_mapping = kwargs.get("registration_mapping", None)
+    if registration_mapping is None and (appending or len(adata_objects) > 1):
+        registration_mapping = soma_io.register_anndatas(
+            experiment_uri=storepath if appending else None,
+            adatas=adata_objects,
+            measurement_name=measurement_name,
+            obs_field_name=obs_id_name,
+            var_field_name=var_id_name,
+            append_obsm_varm=append_obsm_varm,
+            context=ctx,
+        )
+
+    if registration_mapping is not None:
+        n_observations = len(registration_mapping.obs_axis.data)
+    else:  # happens only if not appending and only one adata passed
+        assert len(adata_objects) == 1  # noqa: S101
+        n_observations = adata_objects[0].n_obs
+
+    for adata_obj in adata_objects:
+        soma_io.from_anndata(
+            storepath,
+            adata_obj,
+            measurement_name,
+            context=ctx,
+            obs_id_name=obs_id_name,
+            var_id_name=var_id_name,
+            registration_mapping=registration_mapping,
+            **kwargs,
+        )
+
+    artifact = Artifact(
+        storepath,
+        key=key,
+        description=description,
+        run=run,
+        revises=revises,
+        _is_internal_call=True,
+    )
+    artifact.n_observations = n_observations
+    artifact._accessor = "tiledbsoma"
+
+    return artifact.save()

lamindb/core/storage/_valid_suffixes.py

@@ -1,21 +1,21 @@
-from __future__ import annotations
-
-from lamindb_setup.core.upath import VALID_COMPOSITE_SUFFIXES, VALID_SIMPLE_SUFFIXES
-
-# add new composite suffixes like so
-VALID_COMPOSITE_SUFFIXES.update(
-    {
-        ".vitessce.json",
-        ".ome.zarr",
-    }
-)
-# can do the same for simple valid suffixes
-
-
-class VALID_SUFFIXES:
-    """Valid suffixes."""
-
-    SIMPLE: set[str] = VALID_SIMPLE_SUFFIXES
-    """Simple suffixes."""
-    COMPOSITE: set[str] = VALID_COMPOSITE_SUFFIXES
-    """Composite suffixes."""
+from __future__ import annotations
+
+from lamindb_setup.core.upath import VALID_COMPOSITE_SUFFIXES, VALID_SIMPLE_SUFFIXES
+
+# add new composite suffixes like so
+VALID_COMPOSITE_SUFFIXES.update(
+    {
+        ".vitessce.json",
+        ".ome.zarr",
+    }
+)
+# can do the same for simple valid suffixes
+
+
+class VALID_SUFFIXES:
+    """Valid suffixes."""
+
+    SIMPLE: set[str] = VALID_SIMPLE_SUFFIXES
+    """Simple suffixes."""
+    COMPOSITE: set[str] = VALID_COMPOSITE_SUFFIXES
+    """Composite suffixes."""

lamindb/core/storage/_zarr.py

@@ -1,110 +1,110 @@
-from __future__ import annotations
-
-import warnings
-from typing import TYPE_CHECKING
-
-import scipy.sparse as sparse
-import zarr
-from anndata._io import read_zarr
-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 ._anndata_sizes import _size_elem, _size_raw, size_adata
-
-if TYPE_CHECKING:
-    from anndata import AnnData
-    from lamindb_setup.core.types import UPathStr
-
-
-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:
-    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
-
-
-def write_adata_zarr(
-    adata: AnnData, storepath: UPathStr, callback=None, chunks=None, **dataset_kwargs
-):
-    fs, storepath_str = infer_filesystem(storepath)
-    store = create_mapper(fs, storepath_str, create=True)
-
-    f = zarr.open(store, mode="w")
-
-    adata.strings_to_categoricals()
-    if adata.raw is not None:
-        adata.strings_to_categoricals(adata.raw.var)
-
-    f.attrs.setdefault("encoding-type", "anndata")
-    f.attrs.setdefault("encoding-version", "0.1.0")
-
-    adata_size = None
-    cumulative_val = 0
-
-    def _cb(key_write: str | None = None):
-        nonlocal adata_size
-        nonlocal cumulative_val
-
-        if callback is None:
-            return None
-        if adata_size is None:
-            adata_size = size_adata(adata)
-        if key_write is None:
-            # begin or finish
-            if cumulative_val < adata_size:
-                callback(adata_size, adata_size if cumulative_val > 0 else 0)
-            return None
-
-        elem = getattr(adata, key_write, None)
-        if elem is None:
-            return None
-        elem_size = _size_raw(elem) if key_write == "raw" else _size_elem(elem)
-        if elem_size == 0:
-            return None
-
-        cumulative_val += elem_size
-        callback(adata_size, cumulative_val)
-
-    def _write_elem_cb(f, k, elem, dataset_kwargs):
-        write_elem(f, k, elem, dataset_kwargs=dataset_kwargs)
-        _cb(k)
-
-    _cb(None)
-    with warnings.catch_warnings():
-        warnings.filterwarnings("ignore", category=UserWarning, module="zarr")
-
-        if chunks is not None and not isinstance(adata.X, sparse.spmatrix):
-            _write_elem_cb(
-                f,
-                "X",
-                adata.X,
-                dataset_kwargs=dict(chunks=chunks, **dataset_kwargs),
-            )
-        else:
-            _write_elem_cb(f, "X", adata.X, dataset_kwargs=dataset_kwargs)
-        for elem in ("obs", "var"):
-            _write_elem_cb(f, elem, getattr(adata, elem), dataset_kwargs=dataset_kwargs)
-        for elem in ("obsm", "varm", "obsp", "varp", "layers", "uns"):
-            _write_elem_cb(
-                f, elem, dict(getattr(adata, elem)), dataset_kwargs=dataset_kwargs
-            )
-        _write_elem_cb(f, "raw", adata.raw, dataset_kwargs=dataset_kwargs)
-    # todo: fix size less than total at the end
-    _cb(None)
+from __future__ import annotations
+
+import warnings
+from typing import TYPE_CHECKING
+
+import scipy.sparse as sparse
+import zarr
+from anndata._io import read_zarr
+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 ._anndata_sizes import _size_elem, _size_raw, size_adata
+
+if TYPE_CHECKING:
+    from anndata import AnnData
+    from lamindb_setup.core.types import UPathStr
+
+
+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:
+    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
+
+
+def write_adata_zarr(
+    adata: AnnData, storepath: UPathStr, callback=None, chunks=None, **dataset_kwargs
+):
+    fs, storepath_str = infer_filesystem(storepath)
+    store = create_mapper(fs, storepath_str, create=True)
+
+    f = zarr.open(store, mode="w")
+
+    adata.strings_to_categoricals()
+    if adata.raw is not None:
+        adata.strings_to_categoricals(adata.raw.var)
+
+    f.attrs.setdefault("encoding-type", "anndata")
+    f.attrs.setdefault("encoding-version", "0.1.0")
+
+    adata_size = None
+    cumulative_val = 0
+
+    def _cb(key_write: str | None = None):
+        nonlocal adata_size
+        nonlocal cumulative_val
+
+        if callback is None:
+            return None
+        if adata_size is None:
+            adata_size = size_adata(adata)
+        if key_write is None:
+            # begin or finish
+            if cumulative_val < adata_size:
+                callback(adata_size, adata_size if cumulative_val > 0 else 0)
+            return None
+
+        elem = getattr(adata, key_write, None)
+        if elem is None:
+            return None
+        elem_size = _size_raw(elem) if key_write == "raw" else _size_elem(elem)
+        if elem_size == 0:
+            return None
+
+        cumulative_val += elem_size
+        callback(adata_size, cumulative_val)
+
+    def _write_elem_cb(f, k, elem, dataset_kwargs):
+        write_elem(f, k, elem, dataset_kwargs=dataset_kwargs)
+        _cb(k)
+
+    _cb(None)
+    with warnings.catch_warnings():
+        warnings.filterwarnings("ignore", category=UserWarning, module="zarr")
+
+        if chunks is not None and not isinstance(adata.X, sparse.spmatrix):
+            _write_elem_cb(
+                f,
+                "X",
+                adata.X,
+                dataset_kwargs=dict(chunks=chunks, **dataset_kwargs),
+            )
+        else:
+            _write_elem_cb(f, "X", adata.X, dataset_kwargs=dataset_kwargs)
+        for elem in ("obs", "var"):
+            _write_elem_cb(f, elem, getattr(adata, elem), dataset_kwargs=dataset_kwargs)
+        for elem in ("obsm", "varm", "obsp", "varp", "layers", "uns"):
+            _write_elem_cb(
+                f, elem, dict(getattr(adata, elem)), dataset_kwargs=dataset_kwargs
+            )
+        _write_elem_cb(f, "raw", adata.raw, dataset_kwargs=dataset_kwargs)
+    # todo: fix size less than total at the end
+    _cb(None)