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

lamindb/__init__.py

@@ -1,12 +1,13 @@
 """A data framework for biology.
 
-Tracking notebooks & scripts.
+Tracking notebooks, scripts & functions.
 
 .. autosummary::
    :toctree: .
 
    track
    finish
+   tracked
 
 Registries.
 
@@ -20,7 +21,7 @@ Registries.
    User
    Storage
    Feature
-   FeatureSet
+   Schema
    Param
    Collection
    Project
@@ -33,7 +34,6 @@ Key functionality.
    :toctree: .
 
    connect
-   Curator
    view
    save
 
@@ -44,23 +44,33 @@ Modules and settings.
 
    integrations
    context
+   curators
    settings
+   errors
    setup
    UPath
    base
    core
 
+Backward compatibility.
+
+.. autosummary::
+   :toctree: .
+
+   FeatureSet
+   Curator
+
 """
 
 # denote a release candidate for 0.1.0 with 0.1rc1, 0.1a1, 0.1b1, etc.
-__version__ = "1.0.5"
+__version__ = "1.1.1"
 
 from lamindb_setup._check_setup import InstanceNotSetupError as _InstanceNotSetupError
 from lamindb_setup._check_setup import _check_instance_setup
 from lamindb_setup._connect_instance import connect
 from lamindb_setup.core.upath import UPath
 
-from . import base, setup
+from . import base, errors, setup
 
 
 def __getattr__(name):
@@ -86,10 +96,11 @@ if _check_instance_setup(from_module="lamindb"):
         integrations,
     )
     from ._save import save
+    from ._tracked import tracked
     from ._view import view
     from .core._context import context
     from .core._settings import settings
-    from .curators import Curator
+    from .curators import CatManager as Curator
     from .models import (
         Artifact,
         Collection,

lamindb/_artifact.py

@@ -23,6 +23,8 @@ from lamindb_setup.core.upath import (
     get_stat_file_cloud,
 )
 
+from lamindb._record import _get_record_kwargs
+from lamindb.errors import FieldValidationError
 from lamindb.models import Artifact, FeatureManager, ParamManager, Run, Storage
 
 from ._parents import view_lineage
@@ -32,10 +34,9 @@ from .core._data import (
     describe,
     get_run,
     save_schema_links,
-    save_staged__schemas_m2m,
+    save_staged_feature_sets,
 )
 from .core._settings import settings
-from .core.exceptions import IntegrityError, InvalidArgument
 from .core.loaders import load_to_memory
 from .core.storage import (
     LocalPathClasses,
@@ -44,7 +45,9 @@ from .core.storage import (
     infer_suffix,
     write_to_disk,
 )
+from .core.storage._anndata_accessor import _anndata_n_observations
 from .core.storage._pyarrow_dataset import PYARROW_SUFFIXES
+from .core.storage._tiledbsoma import _soma_n_observations
 from .core.storage.objects import _mudata_is_installed
 from .core.storage.paths import (
     AUTO_KEY_PREFIX,
@@ -58,6 +61,7 @@ from .core.versioning import (
     create_uid,
     message_update_key_in_version_family,
 )
+from .errors import IntegrityError, InvalidArgument
 
 try:
     from .core.storage._zarr import zarr_is_adata
@@ -73,6 +77,7 @@ if TYPE_CHECKING:
     from pyarrow.dataset import Dataset as PyArrowDataset
     from tiledbsoma import Collection as SOMACollection
     from tiledbsoma import Experiment as SOMAExperiment
+    from tiledbsoma import Measurement as SOMAMeasurement
 
     from lamindb.core.storage._backed_access import AnnDataAccessor, BackedAccessor
 
@@ -83,6 +88,7 @@ def process_pathlike(
     using_key: str | None,
     skip_existence_check: bool = False,
 ) -> tuple[Storage, bool]:
+    """Determines the appropriate storage for a given path and whether to use an existing storage key."""
     if not skip_existence_check:
         try:  # check if file exists
             if not filepath.exists():
@@ -112,6 +118,10 @@ def process_pathlike(
                     hf_path.path_in_repo = ""
                     new_root = "hf://" + hf_path.unresolve()
                 else:
+                    if filepath.protocol == "s3":
+                        # check that endpoint_url didn't propagate here
+                        # as a part of the path string
+                        assert "?" not in filepath.path  # noqa: S101
                     new_root = list(filepath.parents)[-1]
                 # do not register remote storage locations on hub if the current instance
                 # is not managed on the hub
@@ -142,6 +152,7 @@ def process_data(
     default_storage: Storage,
     using_key: str | None,
     skip_existence_check: bool = False,
+    is_replace: bool = False,
 ) -> tuple[Any, Path | UPath, str, Storage, bool]:
     """Serialize a data object that's provided as file or in memory."""
     # if not overwritten, data gets stored in default storage
@@ -151,14 +162,24 @@ def process_data(
         data_types = (pd.DataFrame, AnnData, MuData)
     else:
         data_types = (pd.DataFrame, AnnData)  # type:ignore
-
+    if key is not None:
+        key_suffix = extract_suffix_from_path(PurePosixPath(key), arg_name="key")
+        # use suffix as the (adata) format if the format is not provided
+        if isinstance(data, AnnData) and format is None and len(key_suffix) > 0:
+            format = key_suffix[1:]
+    else:
+        key_suffix = None
     if isinstance(data, (str, Path, UPath)):  # UPathStr, spelled out
         access_token = (
             default_storage._access_token
             if hasattr(default_storage, "_access_token")
             else None
         )
-        path = create_path(data, access_token=access_token).resolve()
+        path = create_path(data, access_token=access_token)
+        # we don't resolve http links because they can resolve into a different domain
+        # for example into a temporary url
+        if path.protocol not in {"http", "https"}:
+            path = path.resolve()
         storage, use_existing_storage_key = process_pathlike(
             path,
             default_storage=default_storage,
@@ -170,31 +191,23 @@ def process_data(
     elif isinstance(data, data_types):
         storage = default_storage
         memory_rep = data
-        if key is not None:
-            key_suffix = extract_suffix_from_path(PurePosixPath(key), arg_name="key")
-            # use suffix as the (adata) format if the format is not provided
-            if isinstance(data, AnnData) and format is None and len(key_suffix) > 0:
-                format = key_suffix[1:]
-        else:
-            key_suffix = None
         suffix = infer_suffix(data, format)
-        if key_suffix is not None and key_suffix != suffix:
-            raise InvalidArgument(
-                f"The suffix '{key_suffix}' of the provided key is incorrect, it should"
-                f" be '{suffix}'."
-            )
-        cache_name = f"{provisional_uid}{suffix}"
-        path = settings.cache_dir / cache_name
-        # Alex: I don't understand the line below
-        if path.suffixes == []:
-            path = path.with_suffix(suffix)
-        write_to_disk(data, path)
-        use_existing_storage_key = False
     else:
         raise NotImplementedError(
-            f"Do not know how to create a artifact object from {data}, pass a path"
-            " instead!"
+            f"Do not know how to create a artifact object from {data}, pass a path instead!"
         )
+    if key_suffix is not None and key_suffix != suffix and not is_replace:
+        # consciously omitting a trailing period
+        if isinstance(data, (str, Path, UPath)):
+            message = f"The suffix '{suffix}' of the provided path is inconsistent, it should be '{key_suffix}'"
+        else:
+            message = f"The suffix '{key_suffix}' of the provided key is inconsistent, it should be '{suffix}'"
+        raise InvalidArgument(message)
+    # in case we have an in-memory representation, we need to write it to disk
+    if isinstance(data, data_types):
+        path = settings.cache_dir / f"{provisional_uid}{suffix}"
+        write_to_disk(data, path)
+        use_existing_storage_key = False
     return memory_rep, path, suffix, storage, use_existing_storage_key
 
 
@@ -205,6 +218,7 @@ def get_stat_or_artifact(
     is_replace: bool = False,
     instance: str | None = None,
 ) -> tuple[int, str | None, str | None, int | None, Artifact | None] | Artifact:
+    """Retrieves file statistics or an existing artifact based on the path, hash, and key."""
     n_files = None
     if settings.creation.artifact_skip_size_hash:
         return None, None, None, n_files, None
@@ -248,29 +262,14 @@ def get_stat_or_artifact(
             )
             previous_artifact_version = result[0]
     if artifact_with_same_hash_exists:
-        if settings.creation.artifact_if_hash_exists == "error":
-            msg = f"artifact with same hash exists: {result[0]}"
-            hint = (
-                "💡 you can make this error a warning:\n"
-                "    ln.settings.creation.artifact_if_hash_exists"
-            )
-            raise FileExistsError(f"{msg}\n{hint}")
-        elif settings.creation.artifact_if_hash_exists == "warn_create_new":
-            logger.warning(
-                "creating new Artifact object despite existing artifact with same hash:"
-                f" {result[0]}"
-            )
-            return size, hash, hash_type, n_files, None
-        else:
-            if result[0]._branch_code == -1:
-                raise FileExistsError(
-                    f"You're trying to re-create this artifact in trash: {result[0]}"
-                    "Either permanently delete it with `artifact.delete(permanent=True)` or restore it with `artifact.restore()`"
-                )
-            logger.important(
-                f"returning existing artifact with same hash: {result[0]}; if you intended to query to track this artifact as an input, use: ln.Artifact.get()"
-            )
-            return result[0]
+        message = "found artifact with same hash"
+        if result[0]._branch_code == -1:
+            result[0].restore()
+            message = "restored artifact with same hash from trash"
+        logger.important(
+            f"{message}: {result[0]}; to track this artifact as an input, use: ln.Artifact.get()"
+        )
+        return result[0]
     else:
         return size, hash, hash_type, n_files, previous_artifact_version
 
@@ -326,6 +325,7 @@ def get_artifact_kwargs_from_data(
         default_storage,
         using_key,
         skip_check_exists,
+        is_replace=is_replace,
     )
     stat_or_artifact = get_stat_or_artifact(
         path=path,
@@ -441,7 +441,7 @@ def log_storage_hint(
             root_path = Path(storage.root)  # type: ignore
             if check_path_is_child_of_root(root_path, Path.cwd()):
                 # only display the relative path, not the fully resolved path
-                display_root = root_path.relative_to(Path.cwd())
+                display_root = root_path.relative_to(Path.cwd())  # type: ignore
         hint += f"path in storage '{display_root}'"  # type: ignore
     else:
         hint += "path content will be copied to default storage upon `save()`"
@@ -458,7 +458,7 @@ def data_is_anndata(data: AnnData | UPathStr) -> bool:
         return True
     if isinstance(data, (str, Path, UPath)):
         data_path = UPath(data)
-        if data_path.suffix == ".h5ad":
+        if ".h5ad" in data_path.suffixes:  # ".h5ad.gz" is a valid suffix
             return True
         elif data_path.suffix == ".zarr":
             # ".anndata.zarr" is a valid suffix (core.storage._valid_suffixes)
@@ -480,7 +480,7 @@ def data_is_mudata(data: MuData | UPathStr) -> bool:
         if isinstance(data, MuData):
             return True
     if isinstance(data, (str, Path)):
-        return UPath(data).suffix in {".h5mu"}
+        return UPath(data).suffix == ".h5mu"
     return False
 
 
@@ -506,8 +506,8 @@ def _check_otype_artifact(data: Any, otype: str | None = None):
 
 
 def __init__(artifact: Artifact, *args, **kwargs):
-    artifact.features = FeatureManager(artifact)
-    artifact.params = ParamManager(artifact)
+    artifact.features = FeatureManager(artifact)  # type: ignore
+    artifact.params = ParamManager(artifact)  # type: ignore
     # Below checks for the Django-internal call in from_db()
     # it'd be better if we could avoid this, but not being able to create a Artifact
     # from data with the default constructor renders the central class of the API
@@ -559,9 +559,9 @@ def __init__(artifact: Artifact, *args, **kwargs):
         logger.warning("`type` will be removed soon, please use `kind`")
         kind = kwargs.pop("type")
     if not len(kwargs) == 0:
-        raise ValueError(
-            "Only data, key, run, description, version, revises"
-            f" can be passed, you passed: {kwargs}"
+        valid_keywords = ", ".join([val[0] for val in _get_record_kwargs(Artifact)])
+        raise FieldValidationError(
+            f"Only {valid_keywords} can be passed, you passed: {kwargs}"
         )
     if revises is not None and key is not None and revises.key != key:
         note = message_update_key_in_version_family(
@@ -676,6 +676,7 @@ def __init__(artifact: Artifact, *args, **kwargs):
 def from_df(
     cls,
     df: pd.DataFrame,
+    *,
     key: str | None = None,
     description: str | None = None,
     run: Run | None = None,
@@ -683,7 +684,7 @@ def from_df(
     **kwargs,
 ) -> Artifact:
     """{}"""  # noqa: D415
-    artifact = Artifact(
+    artifact = Artifact(  # type: ignore
         data=df,
         key=key,
         run=run,
@@ -693,6 +694,7 @@ def from_df(
         kind="dataset",
         **kwargs,
     )
+    artifact.n_observations = len(df)
     return artifact
 
 
@@ -701,6 +703,7 @@ def from_df(
 def from_anndata(
     cls,
     adata: AnnData | UPathStr,
+    *,
     key: str | None = None,
     description: str | None = None,
     run: Run | None = None,
@@ -710,7 +713,8 @@ def from_anndata(
     """{}"""  # noqa: D415
     if not data_is_anndata(adata):
         raise ValueError("data has to be an AnnData object or a path to AnnData-like")
-    artifact = Artifact(
+    _anndata_n_observations(adata)
+    artifact = Artifact(  # type: ignore
         data=adata,
         key=key,
         run=run,
@@ -720,6 +724,17 @@ def from_anndata(
         kind="dataset",
         **kwargs,
     )
+    # this is done instead of _anndata_n_observations(adata)
+    # because we need a proper path through create_path for cloud paths
+    # for additional upath options etc that create_path adds
+    obj_for_obs: AnnData | UPath
+    if hasattr(artifact, "_memory_rep") and artifact._memory_rep is not None:
+        obj_for_obs = artifact._memory_rep
+    else:
+        # returns ._local_filepath for local files
+        # and the proper path through create_path for cloud paths
+        obj_for_obs = artifact.path
+    artifact.n_observations = _anndata_n_observations(obj_for_obs)
     return artifact
 
 
@@ -728,6 +743,7 @@ def from_anndata(
 def from_mudata(
     cls,
     mdata: MuData,
+    *,
     key: str | None = None,
     description: str | None = None,
     run: Run | None = None,
@@ -735,7 +751,7 @@ def from_mudata(
     **kwargs,
 ) -> Artifact:
     """{}"""  # noqa: D415
-    artifact = Artifact(
+    artifact = Artifact(  # type: ignore
         data=mdata,
         key=key,
         run=run,
@@ -745,6 +761,38 @@ def from_mudata(
         kind="dataset",
         **kwargs,
     )
+    artifact.n_observations = mdata.n_obs
+    return artifact
+
+
+@classmethod  # type: ignore
+@doc_args(Artifact.from_tiledbsoma.__doc__)
+def from_tiledbsoma(
+    cls,
+    path: UPathStr,
+    *,
+    key: str | None = None,
+    description: str | None = None,
+    run: Run | None = None,
+    revises: Artifact | None = None,
+    **kwargs,
+) -> Artifact:
+    """{}"""  # noqa: D415
+    if UPath(path).suffix != ".tiledbsoma":
+        raise ValueError(
+            "A tiledbsoma store should have .tiledbsoma suffix to be registered."
+        )
+    artifact = Artifact(  # type: ignore
+        data=path,
+        key=key,
+        run=run,
+        description=description,
+        revises=revises,
+        otype="tiledbsoma",
+        kind="dataset",
+        **kwargs,
+    )
+    artifact.n_observations = _soma_n_observations(artifact.path)
     return artifact
 
 
@@ -753,8 +801,8 @@ def from_mudata(
 def from_dir(
     cls,
     path: UPathStr,
-    key: str | None = None,
     *,
+    key: str | None = None,
     run: Run | None = None,
 ) -> list[Artifact]:
     """{}"""  # noqa: D415
@@ -931,22 +979,34 @@ inconsistent_state_msg = (
 
 # docstring handled through attach_func_to_class_method
 def open(
-    self, mode: str = "r", is_run_input: bool | None = None
+    self, mode: str = "r", is_run_input: bool | None = None, **kwargs
 ) -> (
-    AnnDataAccessor | BackedAccessor | SOMACollection | SOMAExperiment | PyArrowDataset
+    AnnDataAccessor
+    | BackedAccessor
+    | SOMACollection
+    | SOMAExperiment
+    | SOMAMeasurement
+    | PyArrowDataset
 ):
     if self._overwrite_versions and not self.is_latest:
         raise ValueError(inconsistent_state_msg)
+    # all hdf5 suffixes including gzipped
+    h5_suffixes = [".h5", ".hdf5", ".h5ad"]
+    h5_suffixes += [s + ".gz" for s in h5_suffixes]
     # ignore empty suffix for now
     suffixes = (
-        "",
-        ".h5",
-        ".hdf5",
-        ".h5ad",
-        ".zarr",
-        ".anndata.zarr",
-        ".tiledbsoma",
-    ) + PYARROW_SUFFIXES
+        (
+            "",
+            ".zarr",
+            ".anndata.zarr",
+            ".tiledbsoma",
+        )
+        + tuple(h5_suffixes)
+        + PYARROW_SUFFIXES
+        + tuple(
+            s + ".gz" for s in PYARROW_SUFFIXES
+        )  # this doesn't work for externally gzipped files, REMOVE LATER
+    )
     if self.suffix not in suffixes:
         raise ValueError(
             "Artifact should have a zarr, h5, tiledbsoma object"
@@ -964,16 +1024,36 @@ def open(
     using_key = settings._using_key
     filepath, cache_key = filepath_cache_key_from_artifact(self, using_key=using_key)
     is_tiledbsoma_w = (
-        filepath.name == "soma" or filepath.suffix == ".tiledbsoma"
+        filepath.name == "soma" or self.suffix == ".tiledbsoma"
     ) and mode == "w"
     # consider the case where an object is already locally cached
     localpath = setup_settings.paths.cloud_to_local_no_update(
         filepath, cache_key=cache_key
     )
-    if not is_tiledbsoma_w and localpath.exists():
-        access = backed_access(localpath, mode, using_key)
+    if is_tiledbsoma_w:
+        open_cache = False
     else:
-        access = backed_access(filepath, mode, using_key)
+        open_cache = not isinstance(
+            filepath, LocalPathClasses
+        ) and not filepath.synchronize(localpath, just_check=True)
+    if open_cache:
+        try:
+            access = backed_access(localpath, mode, using_key, **kwargs)
+        except Exception as e:
+            if isinstance(filepath, LocalPathClasses):
+                raise e
+            logger.warning(
+                f"The cache might be corrupted: {e}. Trying to open directly."
+            )
+            access = backed_access(filepath, mode, using_key, **kwargs)
+            # happens only if backed_access has been successful
+            # delete the corrupted cache
+            if localpath.is_dir():
+                shutil.rmtree(localpath)
+            else:
+                localpath.unlink(missing_ok=True)
+    else:
+        access = backed_access(filepath, mode, using_key, **kwargs)
         if is_tiledbsoma_w:
 
             def finalize():
@@ -1013,10 +1093,10 @@ def _synchronize_cleanup_on_error(
             cache_path = setup_settings.paths.cloud_to_local_no_update(
                 filepath, cache_key=cache_key
             )
-            if cache_path.is_file():
-                cache_path.unlink(missing_ok=True)
-            elif cache_path.is_dir():
+            if cache_path.is_dir():
                 shutil.rmtree(cache_path)
+            else:
+                cache_path.unlink(missing_ok=True)
         raise e
     return cache_path
 
@@ -1033,8 +1113,24 @@ def load(self, is_run_input: bool | None = None, **kwargs) -> Any:
             self, using_key=settings._using_key
         )
         cache_path = _synchronize_cleanup_on_error(filepath, cache_key=cache_key)
-        # cache_path is local so doesn't trigger any sync in load_to_memory
-        access_memory = load_to_memory(cache_path, **kwargs)
+        try:
+            # cache_path is local so doesn't trigger any sync in load_to_memory
+            access_memory = load_to_memory(cache_path, **kwargs)
+        except Exception as e:
+            # just raise the exception if the original path is local
+            if isinstance(filepath, LocalPathClasses):
+                raise e
+            logger.warning(
+                f"The cache might be corrupted: {e}. Retrying to synchronize."
+            )
+            # delete the existing cache
+            if cache_path.is_dir():
+                shutil.rmtree(cache_path)
+            else:
+                cache_path.unlink(missing_ok=True)
+            # download again and try to load into memory
+            cache_path = _synchronize_cleanup_on_error(filepath, cache_key=cache_key)
+            access_memory = load_to_memory(cache_path, **kwargs)
     # only call if load is successfull
     _track_run_input(self, is_run_input)
     return access_memory
@@ -1154,6 +1250,7 @@ def _delete_skip_storage(artifact, *args, **kwargs) -> None:
 def save(self, upload: bool | None = None, **kwargs) -> Artifact:
     state_was_adding = self._state.adding
     print_progress = kwargs.pop("print_progress", True)
+    store_kwargs = kwargs.pop("store_kwargs", {})  # kwargs for .upload_from in the end
     access_token = kwargs.pop("access_token", None)
     local_path = None
     if upload and setup_settings.instance.keep_artifacts_local:
@@ -1174,15 +1271,31 @@ def save(self, upload: bool | None = None, **kwargs) -> Artifact:
     using_key = None
     if "using" in kwargs:
         using_key = kwargs["using"]
-    exception = check_and_attempt_upload(
-        self, using_key, access_token=access_token, print_progress=print_progress
+    exception_upload = check_and_attempt_upload(
+        self,
+        using_key,
+        access_token=access_token,
+        print_progress=print_progress,
+        **store_kwargs,
     )
-    if exception is not None:
+    if exception_upload is not None:
+        # we do not want to raise file not found on cleanup if upload of a file failed
+        # often it is ACID in the filesystem itself
+        # for example, s3 won't have the failed file, so just skip the delete in this case
+        raise_file_not_found_error = False
         self._delete_skip_storage()
-        raise RuntimeError(exception)
-    exception = check_and_attempt_clearing(self, using_key)
-    if exception is not None:
-        raise RuntimeError(exception)
+    else:
+        # this is the case when it is cleaned on .replace
+        raise_file_not_found_error = True
+    # this is triggered by an exception in check_and_attempt_upload or by replace.
+    exception_clear = check_and_attempt_clearing(
+        self, raise_file_not_found_error=raise_file_not_found_error, using_key=using_key
+    )
+    if exception_upload is not None:
+        raise RuntimeError(exception_upload)
+    if exception_clear is not None:
+        raise RuntimeError(exception_clear)
+    # this is only for keep_artifacts_local
     if local_path is not None and not state_was_adding:
         # only move the local artifact to cache if it was not newly created
         local_path_cache = ln_setup.settings.cache_dir / local_path.name
@@ -1197,7 +1310,7 @@ def save(self, upload: bool | None = None, **kwargs) -> Artifact:
 
 
 def _save_skip_storage(file, **kwargs) -> None:
-    save_staged__schemas_m2m(file)
+    save_staged_feature_sets(file)
     super(Artifact, file).save(**kwargs)
     save_schema_links(file)
 
@@ -1233,6 +1346,7 @@ METHOD_NAMES = [
     "from_anndata",
     "from_df",
     "from_mudata",
+    "from_tiledbsoma",
     "open",
     "cache",
     "load",
@@ -1256,6 +1370,7 @@ for name in METHOD_NAMES:
     attach_func_to_class_method(name, Artifact, globals())
 
 # privates currently dealt with separately
+# mypy: ignore-errors
 Artifact._delete_skip_storage = _delete_skip_storage
 Artifact._save_skip_storage = _save_skip_storage
 Artifact._cache_path = _cache_path