Skip to content

Commit 0c5078f

Browse files
authored
🚸 Enable passing --branch and --space to lamin save (#2819)
1 parent f4125cc commit 0c5078f

File tree

15 files changed

+621
-76
lines changed

15 files changed

+621
-76
lines changed

‎lamindb/core/_context.py

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@
3838
from lamindb_setup.core.types import UPathStr
3939

4040
from lamindb.base.types import TransformType
41-
from lamindb.models import Project, Space
41+
from lamindb.models import Branch, Project, Space
4242

4343
is_run_from_ipython = getattr(builtins, "__IPYTHON__", False)
4444

@@ -193,6 +193,7 @@ def __init__(self):
193193
"""A local path to the script or notebook that's running."""
194194
self._project: Project | None = None
195195
self._space: Space | None = None
196+
self._branch: Branch | None = None
196197
self._logging_message_track: str = ""
197198
self._logging_message_imports: str = ""
198199
self._stream_tracker: LogStreamTracker = LogStreamTracker()
@@ -247,9 +248,14 @@ def project(self) -> Project | None:
247248

248249
@property
249250
def space(self) -> Space | None:
250-
"""The space in which entities are created during the run."""
251+
"""The space in which artifacts, collections, transforms, and runs are saved during the run."""
251252
return self._space
252253

254+
@property
255+
def branch(self) -> Branch | None:
256+
"""The branch on which entities are created during the run."""
257+
return self._branch
258+
253259
@property
254260
def run(self) -> Run | None:
255261
"""Managed run of context."""
@@ -261,6 +267,7 @@ def _track(
261267
*,
262268
project: str | Project | None = None,
263269
space: str | Space | None = None,
270+
branch: str | Branch | None = None,
264271
params: dict | None = None,
265272
new_run: bool | None = None,
266273
path: str | None = None,
@@ -273,11 +280,11 @@ def _track(
273280
274281
Args:
275282
transform: A transform (stem) `uid` (or record). If `None`, auto-creates a `transform` with its `uid`.
276-
project: A project, its `name` or `uid` for labeling entities created during the run.
277-
space: A restricted space, its `name` or `uid` for creating sensitive entities are created during the run.
278-
The default is the common `"All"` space that every LaminDB instance has.
279-
The `space` argument doesn't affect `Storage`, `ULabel`, `Feature`, `Schema`, `Param` and bionty entities as these provide structure that should typically be commonly accessible.
283+
project: A project (or its `name` or `uid`) for labeling entities.
284+
space: A restricted space (or its `name` or `uid`) in which to store artifacts, collections, transforms, and runs.
285+
Default: the `"All"` space.
280286
If you want to manually move entities to a different space, set the `.space` field (:doc:`docs:access`).
287+
branch: A branch (or its `name` or `uid`) on which to store records.
281288
params: A dictionary of parameters to track for the run.
282289
new_run: If `False`, loads the latest run of transform
283290
(default notebook), if `True`, creates new run (default non-notebook).
@@ -299,7 +306,7 @@ def _track(
299306
300307
More examples: :doc:`/track`
301308
"""
302-
from lamindb.models import Project, Space
309+
from lamindb.models import Branch, Project, Space
303310

304311
instance_settings = ln_setup.settings.instance
305312
# similar logic here: https://github.com/laminlabs/lamindb/pull/2527
@@ -337,6 +344,21 @@ def _track(
337344
f"Space '{space}', please check on the hub UI whether you have the correct `uid` or `name`."
338345
)
339346
self._space = space_record
347+
if branch is not None:
348+
if isinstance(branch, Branch):
349+
assert branch._state.adding is False, ( # noqa: S101
350+
"Branch must be saved before passing it to track()"
351+
)
352+
branch_record = branch
353+
else:
354+
branch_record = Branch.filter(
355+
Q(name=branch) | Q(uid=branch)
356+
).one_or_none()
357+
if branch_record is None:
358+
raise InvalidArgument(
359+
f"Space '{branch}', please check on the hub UI whether you have the correct `uid` or `name`."
360+
)
361+
self._branch = branch_record
340362
self._logging_message_track = ""
341363
self._logging_message_imports = ""
342364
if transform is not None and isinstance(transform, str):

0 commit comments

Comments
 (0)