lamindb 1.3.2__py3-none-any.whl → 1.5.0__py3-none-any.whl

lamindb/core/_context.py

@@ -2,6 +2,7 @@ from __future__ import annotations
 
 import builtins
 import hashlib
+import os
 import signal
 import sys
 import threading
@@ -22,7 +23,6 @@ from lamindb.models import Run, Transform, format_field_value
 
 from ..core._settings import settings
 from ..errors import (
-    InconsistentKey,
     InvalidArgument,
     TrackNotCalled,
     UpdateContext,
@@ -30,7 +30,6 @@ from ..errors import (
 from ..models._is_versioned import bump_version as bump_version_function
 from ..models._is_versioned import (
     increment_base62,
-    message_update_key_in_version_family,
 )
 from ._sync_git import get_transform_reference_from_git_repo
 from ._track_environment import track_environment
@@ -39,7 +38,7 @@ if TYPE_CHECKING:
     from lamindb_setup.core.types import UPathStr
 
     from lamindb.base.types import TransformType
-    from lamindb.models import Project
+    from lamindb.models import Project, Space
 
 is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
 
@@ -55,19 +54,21 @@ def get_uid_ext(version: str) -> str:
     return encodebytes(hashlib.md5(version.encode()).digest())[:4]  # noqa: S324
 
 
-def get_notebook_path() -> Path:
+def get_notebook_path() -> tuple[Path, str]:
     from nbproject.dev._jupyter_communicate import (
         notebook_path as get_notebook_path,
     )
 
     path = None
     try:
-        path = get_notebook_path()
-    except Exception:
-        raise RuntimeError(msg_path_failed) from None
+        path, env = get_notebook_path(return_env=True)
+    except ValueError as ve:
+        raise ve
+    except Exception as error:
+        raise RuntimeError(msg_path_failed) from error
     if path is None:
         raise RuntimeError(msg_path_failed) from None
-    return Path(path)
+    return Path(path), env
 
 
 # from https://stackoverflow.com/questions/61901628
@@ -179,20 +180,7 @@ class LogStreamTracker:
 class Context:
     """Run context.
 
-    Enables convenient data lineage tracking by managing a transform & run
-    upon :meth:`~lamindb.core.Context.track` & :meth:`~lamindb.core.Context.finish`.
-
-    Guide: :doc:`/track`
-
-    Examples:
-
-        Is typically used via the global :class:`~lamindb.context` object via `ln.track()` and `ln.finish()`:
-
-        >>> import lamindb as ln
-        >>> ln.track()
-        >>> # do things
-        >>> ln.finish()
-
+    Is the book keeper for :meth:`~lamindb.core.Context.track`.
     """
 
     def __init__(self):
@@ -202,12 +190,14 @@ class Context:
         self._transform: Transform | None = None
         self._run: Run | None = None
         self._path: Path | None = None
-        """A local path to the script that's running."""
+        """A local path to the script or notebook that's running."""
         self._project: Project | None = None
+        self._space: Space | None = None
         self._logging_message_track: str = ""
         self._logging_message_imports: str = ""
         self._stream_tracker: LogStreamTracker = LogStreamTracker()
         self._is_finish_retry: bool = False
+        self._notebook_runner: str | None = None
 
     @property
     def transform(self) -> Transform | None:
@@ -255,33 +245,38 @@ class Context:
         """Project to label entities created during the run."""
         return self._project
 
+    @property
+    def space(self) -> Space | None:
+        """The space in which entities are created during the run."""
+        return self._space
+
     @property
     def run(self) -> Run | None:
         """Managed run of context."""
         return self._run
 
-    def track(
+    def _track(
         self,
         transform: str | Transform | None = None,
         *,
         project: str | None = None,
+        space: str | None = None,
         params: dict | None = None,
         new_run: bool | None = None,
         path: str | None = None,
     ) -> None:
-        """Track a global run of your Python session.
+        """Track a run of your notebook or script.
 
-        - sets :attr:`~lamindb.core.Context.transform` &
-          :attr:`~lamindb.core.Context.run` by creating or loading `Transform` &
-          `Run` records
-        - saves Python environment as a `requirements.txt` file: `run.environment`
+        Populates the global run :class:`~lamindb.context` by managing `Transform` & `Run` records and caching the compute environment.
 
-        If :attr:`~lamindb.core.Settings.sync_git_repo` is set, checks whether a
-        script-like transform exists in a git repository and links it.
+        If :attr:`~lamindb.core.Settings.sync_git_repo` is set, checks whether a script-like transform exists in a git repository and links it.
 
         Args:
-            transform: A transform `uid` or record. If `None`, creates a `uid`.
+            transform: A transform (stem) `uid` (or record). If `None`, auto-creates a `transform` with its `uid`.
             project: A project `name` or `uid` for labeling entities created during the run.
+            space: A space `name` or `uid` to identify where potentially sensitive entities are created during the run.
+                This doesn't affect `Storage`, `ULabel`, `Feature`, `Schema`, `Param` and bionty entities as these provide mere structure that should typically be commonly accessible.
+                If you want to manually move entities to a different space, set the `.space` field (:doc:`docs:access`).
             params: A dictionary of parameters to track for the run.
             new_run: If `False`, loads the latest run of transform
                 (default notebook), if `True`, creates new run (default non-notebook).
@@ -290,16 +285,20 @@ class Context:
 
         Examples:
 
-            To track the run of a notebook or script, call:
+            To track the run of a notebook or script, call::
 
-            >>> ln.track()
+                ln.track()
+                #> → created Transform('Onv04I53OgtT0000'), started new Run('dpSfd7Ds...') at 2025-04-25 11:00:03 UTC
+                #> • recommendation: to identify the notebook across renames, pass the uid: ln.track("Onv04I53OgtT")
 
-            If you want to ensure a single version history across renames of the notebook or script, pass the auto-generated `uid` that you'll find in the logs:
+            Ensure one version history across file renames::
 
-            >>> ln.track("Onv04I53OgtT0000")  # example uid, the last four characters encode the version of the transform
+                ln.track("Onv04I53OgtT")
+                #> → created Transform('Onv04I53OgtT0000'), started new Run('dpSfd7Ds...') at 2025-04-25 11:00:03 UTC
 
+            More examples: :doc:`/track`
         """
-        from lamindb.models import Project
+        from lamindb.models import Project, Space
 
         instance_settings = ln_setup.settings.instance
         # similar logic here: https://github.com/laminlabs/lamindb/pull/2527
@@ -307,6 +306,8 @@ class Context:
         if instance_settings.dialect == "postgresql" and "read" in instance_settings.db:
             logger.warning("skipping track(), connected in read-only mode")
             return None
+        if project is None:
+            project = os.environ.get("LAMIN_CURRENT_PROJECT")
         if project is not None:
             project_record = Project.filter(
                 Q(name=project) | Q(uid=project)
@@ -316,11 +317,21 @@ class Context:
                     f"Project '{project}' not found, either create it with `ln.Project(name='...').save()` or fix typos."
                 )
             self._project = project_record
+        if space is not None:
+            space_record = Space.filter(Q(name=space) | Q(uid=space)).one_or_none()
+            if space_record is None:
+                raise InvalidArgument(
+                    f"Space '{space}', please check on the hub UI whether you have the correct `uid` or `name`."
+                )
+            self._space = space_record
         self._logging_message_track = ""
         self._logging_message_imports = ""
         if transform is not None and isinstance(transform, str):
             self.uid = transform
             transform = None
+            uid_was_none = False
+        else:
+            uid_was_none = True
         self._path = None
         if transform is None:
             description = None
@@ -367,7 +378,14 @@ class Context:
             self._transform = transform_exists
 
         if new_run is None:  # for notebooks, default to loading latest runs
-            new_run = False if self._transform.type == "notebook" else True  # type: ignore
+            new_run = (
+                False
+                if (
+                    self._transform.type == "notebook"
+                    and self._notebook_runner != "nbconvert"
+                )
+                else True
+            )  # type: ignore
 
         run = None
         if not new_run:  # try loading latest run by same user
@@ -414,6 +432,22 @@ class Context:
         logger.important(self._logging_message_track)
         if self._logging_message_imports:
             logger.important(self._logging_message_imports)
+        if uid_was_none:
+            notebook_or_script = (
+                "notebook" if self._transform.type == "notebook" else "script"
+            )
+            r_or_python = "."
+            if self._path is not None:
+                r_or_python = "." if self._path.suffix in {".py", ".ipynb"} else "$"
+            project_str = f', project="{project}"' if project is not None else ""
+            space_str = f', space="{space}"' if space is not None else ""
+            params_str = (
+                ", params={...}" if params is not None else ""
+            )  # do not put the values because typically parameterized by user
+            kwargs_str = f"{project_str}{space_str}{params_str}"
+            logger.important_hint(
+                f'recommendation: to identify the {notebook_or_script} across renames, pass the uid: ln{r_or_python}track("{self.transform.uid[:-4]}"{kwargs_str})'
+            )
 
     def _track_source_code(
         self,
@@ -455,7 +489,7 @@ class Context:
         path_str: str | None,
     ) -> tuple[Path, str | None]:
         if path_str is None:
-            path = get_notebook_path()
+            path, self._notebook_runner = get_notebook_path()
         else:
             path = Path(path_str)
         description = None
@@ -487,6 +521,53 @@ class Context:
                 pass
         return path, description
 
+    def _process_aux_transform(
+        self,
+        aux_transform: Transform,
+        transform_hash: str,
+    ) -> tuple[str, Transform | None, str]:
+        # first part of the if condition: no version bump, second part: version bump
+        message = ""
+        if (
+            # if a user hasn't yet saved the transform source code AND is the same user
+            (
+                aux_transform.source_code is None
+                and aux_transform.created_by_id == ln_setup.settings.user.id
+            )
+            # if the transform source code is unchanged
+            # if aux_transform.type == "notebook", we anticipate the user makes changes to the notebook source code
+            # in an interactive session, hence we *pro-actively bump* the version number by setting `revises` / 'nbconvert' execution is NOT interactive
+            # in the second part of the if condition even though the source code is unchanged at point of running track()
+            or (
+                aux_transform.hash == transform_hash
+                and (
+                    aux_transform.type != "notebook"
+                    or self._notebook_runner == "nbconvert"
+                )
+            )
+        ):
+            uid = aux_transform.uid
+            return uid, aux_transform, message
+        else:
+            uid = f"{aux_transform.uid[:-4]}{increment_base62(aux_transform.uid[-4:])}"
+            message = (
+                f"found {aux_transform.type} {aux_transform.key}, making new version"
+            )
+            if (
+                aux_transform.hash == transform_hash
+                and aux_transform.type == "notebook"
+            ):
+                message += " -- anticipating changes"
+            elif aux_transform.hash != transform_hash:
+                message += (
+                    ""  # could log "source code changed", but this seems too much
+                )
+            elif aux_transform.created_by_id != ln_setup.settings.user.id:
+                message += (
+                    f" -- {aux_transform.created_by.handle} already works on this draft"
+                )
+            return uid, None, message
+
     def _create_or_load_transform(
         self,
         *,
@@ -495,21 +576,22 @@ class Context:
         transform_ref_type: str | None = None,
         transform_type: TransformType = None,
     ):
-        def get_key_clashing_message(transform: Transform, key: str) -> str:
-            update_key_note = message_update_key_in_version_family(
-                suid=transform.stem_uid,
-                existing_key=transform.key,
-                new_key=key,
-                registry="Transform",
-            )
-            return (
-                f'Filepath "{key}" clashes with the existing key "{transform.key}" for uid "{transform.uid[:-4]}...."\n\nEither init a new transform with a new uid:\n\n'
-                f'ln.track("{ids.base62_12()}0000")\n\n{update_key_note}'
-            )
+        from .._finish import notebook_to_script
 
-        revises = None
-        # the user did not pass the uid
-        if self.uid is None:
+        if not self._path.suffix == ".ipynb":
+            transform_hash, _ = hash_file(self._path)
+        else:
+            # need to convert to stripped py:percent format for hashing
+            source_code_path = ln_setup.settings.cache_dir / self._path.name.replace(
+                ".ipynb", ".py"
+            )
+            notebook_to_script(description, self._path, source_code_path)
+            transform_hash, _ = hash_file(source_code_path)
+        # see whether we find a transform with the exact same hash
+        aux_transform = Transform.filter(hash=transform_hash).one_or_none()
+        # if the user did not pass a uid and there is no matching aux_transform
+        # need to search for the transform based on the filename
+        if self.uid is None and aux_transform is None:
 
             class SlashCount(Func):
                 template = "LENGTH(%(expressions)s) - LENGTH(REPLACE(%(expressions)s, '/', ''))"
@@ -524,47 +606,15 @@ class Context:
             uid = f"{base62_12()}0000"
             key = self._path.name
             target_transform = None
-            hash, _ = hash_file(self._path)
             if len(transforms) != 0:
                 message = ""
                 found_key = False
                 for aux_transform in transforms:
                     if aux_transform.key in self._path.as_posix():
                         key = aux_transform.key
-                        if (
-                            # has to be the same user
-                            aux_transform.created_by_id == ln_setup.settings.user.id
-                            and (
-                                # if the transform source code wasn't yet saved
-                                aux_transform.source_code is None
-                                # if the transform source code is unchanged
-                                # if aux_transform.type == "notebook", we anticipate the user makes changes to the notebook source code
-                                # in an interactive session, hence we *pro-actively bump* the version number by setting `revises`
-                                # in the second part of the if condition even though the source code is unchanged at point of running track()
-                                or (
-                                    aux_transform.hash == hash
-                                    and aux_transform.type != "notebook"
-                                )
-                            )
-                        ):
-                            uid = aux_transform.uid
-                            target_transform = aux_transform
-                        else:
-                            uid = f"{aux_transform.uid[:-4]}{increment_base62(aux_transform.uid[-4:])}"
-                            message = f"there already is a {aux_transform.type} with `key` '{aux_transform.key}'"
-                            if (
-                                aux_transform.hash == hash
-                                and aux_transform.type == "notebook"
-                            ):
-                                message += " -- anticipating changes"
-                            elif aux_transform.hash != hash:
-                                message += ""  # could log "source code changed", but this seems too much
-                            elif (
-                                aux_transform.created_by_id != ln_setup.settings.user.id
-                            ):
-                                message += f" -- {aux_transform.created_by.handle} already works on this draft"
-                            message += f", creating new version '{uid}'"
-                            revises = aux_transform
+                        uid, target_transform, message = self._process_aux_transform(
+                            aux_transform, transform_hash
+                        )
                         found_key = True
                         break
                 if not found_key:
@@ -580,22 +630,48 @@ class Context:
                     logger.important(message)
             self.uid, transform = uid, target_transform
         # the user did pass the uid
-        else:
+        elif self.uid is not None and len(self.uid) == 16:
             transform = Transform.filter(uid=self.uid).one_or_none()
             if transform is not None:
                 if transform.key not in self._path.as_posix():
                     n_parts = len(Path(transform.key).parts)
-                    last_path_elements = (
+                    (
                         Path(*self._path.parts[-n_parts:]).as_posix()
                         if n_parts > 0
                         else ""
                     )
-                    raise UpdateContext(
-                        get_key_clashing_message(transform, last_path_elements)
+                    key = self._path.name
+                else:
+                    key = transform.key  # type: ignore
+            else:
+                key = self._path.name
+        else:
+            if self.uid is not None:
+                assert len(self.uid) == 12, (  # noqa: S101
+                    "uid must be 12 (stem) or 16 (full) characters long"
+                )
+                aux_transform = (
+                    Transform.filter(uid__startswith=self.uid)
+                    .order_by("-created_at")
+                    .first()
+                )
+            if aux_transform is not None:
+                if aux_transform.key.endswith(self._path.name):
+                    key = aux_transform.key
+                else:
+                    key = "/".join(
+                        aux_transform.key.split("/")[:-1] + [self._path.name]
                     )
-                key = transform.key  # type: ignore
+                uid, target_transform, message = self._process_aux_transform(
+                    aux_transform, transform_hash
+                )
+                if message != "":
+                    logger.important(message)
             else:
+                uid = f"{self.uid}0000" if self.uid is not None else None
+                target_transform = None
                 key = self._path.name
+            self.uid, transform = uid, target_transform
         if self.version is not None:
             # test inconsistent version passed
             if (
@@ -607,39 +683,28 @@ class Context:
                     f"✗ please pass consistent version: ln.context.version = '{transform.version}'"  # type: ignore
                 )
             # test whether version was already used for another member of the family
-            suid, vuid = (self.uid[:-4], self.uid[-4:])
-            transform = Transform.filter(
-                uid__startswith=suid, version=self.version
-            ).one_or_none()
-            if transform is not None and vuid != transform.uid[-4:]:
-                better_version = bump_version_function(self.version)
-                raise SystemExit(
-                    f"✗ version '{self.version}' is already taken by Transform('{transform.uid}'); please set another version, e.g., ln.context.version = '{better_version}'"
-                )
+            if self.uid is not None and len(self.uid) == 16:
+                suid, vuid = (self.uid[:-4], self.uid[-4:])
+                transform = Transform.filter(
+                    uid__startswith=suid, version=self.version
+                ).one_or_none()
+                if transform is not None and vuid != transform.uid[-4:]:
+                    better_version = bump_version_function(self.version)
+                    raise SystemExit(
+                        f"✗ version '{self.version}' is already taken by Transform('{transform.uid}'); please set another version, e.g., ln.context.version = '{better_version}'"
+                    )
         # make a new transform record
         if transform is None:
             assert key is not None  # noqa: S101
-            raise_update_context = False
-            try:
-                transform = Transform(  # type: ignore
-                    uid=self.uid,
-                    version=self.version,
-                    description=description,
-                    key=key,
-                    reference=transform_ref,
-                    reference_type=transform_ref_type,
-                    type=transform_type,
-                ).save()
-            except InconsistentKey:
-                raise_update_context = True
-            if raise_update_context:
-                if revises is None:
-                    revises = (
-                        Transform.filter(uid__startswith=self.uid[:-4], is_latest=True)
-                        .order_by("-created_at")
-                        .first()
-                    )
-                raise UpdateContext(get_key_clashing_message(revises, key))
+            transform = Transform(  # type: ignore
+                uid=self.uid,
+                version=self.version,
+                description=description,
+                key=key,
+                reference=transform_ref,
+                reference_type=transform_ref_type,
+                type=transform_type,
+            ).save()
             self._logging_message_track += f"created Transform('{transform.uid}')"
         else:
             uid = transform.uid
@@ -647,7 +712,11 @@ class Context:
             transform_was_saved = transform.source_code is not None
             # check whether the transform.key is consistent
             if transform.key != key:
-                raise UpdateContext(get_key_clashing_message(transform, key))
+                self._logging_message_track += (
+                    f"renaming transform {transform.key} to {key}"
+                )
+                transform.key = key
+                transform.save()
             elif transform.description != description and description is not None:
                 transform.description = description
                 transform.save()
@@ -659,18 +728,20 @@ class Context:
                 and not transform_was_saved
             ):
                 raise UpdateContext(
-                    f'{transform.created_by.name} ({transform.created_by.handle}) already works on this draft {transform.type}.\n\nPlease create a revision via `ln.track("{uid[:-4]}{increment_base62(uid[-4:])}")` or a new transform with a *different* filedescription and `ln.track("{ids.base62_12()}0000")`.'
+                    f'{transform.created_by.name} ({transform.created_by.handle}) already works on this draft {transform.type}.\n\nPlease create a revision via `ln.track("{uid[:-4]}{increment_base62(uid[-4:])}")` or a new transform with a *different* key and `ln.track("{ids.base62_12()}0000")`.'
                 )
             # check whether transform source code was already saved
             if transform_was_saved:
                 bump_revision = False
-                if transform.type == "notebook":
+                if (
+                    transform.type == "notebook"
+                    and self._notebook_runner != "nbconvert"
+                ):
                     # we anticipate the user makes changes to the notebook source code
                     # in an interactive session, hence we pro-actively bump the version number
                     bump_revision = True
                 else:
-                    hash, _ = hash_file(self._path)  # ignore hash_type for now
-                    if hash != transform.hash:
+                    if transform_hash != transform.hash:
                         bump_revision = True
                     else:
                         self._logging_message_track += (
@@ -679,7 +750,10 @@ class Context:
                 if bump_revision:
                     change_type = (
                         "re-running notebook with already-saved source code"
-                        if transform.type == "notebook"
+                        if (
+                            transform.type == "notebook"
+                            and self._notebook_runner != "nbconvert"
+                        )
                         else "source code changed"
                     )
                     raise UpdateContext(
@@ -689,11 +763,11 @@ class Context:
                 self._logging_message_track += f"loaded Transform('{transform.uid}')"
         self._transform = transform
 
-    def finish(self, ignore_non_consecutive: None | bool = None) -> None:
-        """Finish a tracked run.
+    def _finish(self, ignore_non_consecutive: None | bool = None) -> None:
+        """Finish the run and write a run report.
 
         - writes a timestamp: `run.finished_at`
-        - saves the source code: `transform.source_code`
+        - saves the source code if it is not yet saved: `transform.source_code`
         - saves a run report: `run.report`
 
         When called in the last cell of a notebook:
@@ -737,6 +811,7 @@ class Context:
             finished_at=True,
             ignore_non_consecutive=ignore_non_consecutive,
             is_retry=self._is_finish_retry,
+            notebook_runner=self._notebook_runner,
         )
         if return_code == "retry":
             self._is_finish_retry = True
@@ -751,5 +826,13 @@ class Context:
         self._version = None
         self._description = None
 
+    @deprecated("ln.track()")
+    def track(self, *args, **kwargs):
+        return self._track(*args, **kwargs)
+
+    @deprecated("ln.finish()")
+    def finish(self, *args, **kwargs):
+        return self._finish(*args, **kwargs)
+
 
-context = Context()
+context: Context = Context()