♻️ Simplify ln.track()

lamin_cli/_save.py

@@ -14,6 +14,7 @@ def get_stem_uid_and_version_from_file(
         content = file.read()
 
     if file_path.suffix == ".py":
+        track_pattern = re.compile(r'ln\.track\(\s*(?:uid\s*=\s*)?["\']([^"\']+)["\']')
         uid_pattern = re.compile(r'\.context\.uid\s*=\s*["\']([^"\']+)["\']')
         stem_uid_pattern = re.compile(
             r'\.transform\.stem_uid\s*=\s*["\']([^"\']+)["\']'
@@ -31,7 +32,10 @@ def get_stem_uid_and_version_from_file(
         raise ValueError("Only .py and .ipynb files are supported.")
 
     # Search for matches in the entire file content
-    uid_match = uid_pattern.search(content)
+    uid_match = track_pattern.search(content)
+    uid = uid_match.group(1) if uid_match else None
+    if uid is None:
+        uid_match = uid_pattern.search(content)
     stem_uid_match = stem_uid_pattern.search(content)
     version_match = version_pattern.search(content)
 
@@ -42,8 +46,8 @@ def get_stem_uid_and_version_from_file(
 
     if uid is None and (stem_uid is None or version is None):
         raise SystemExit(
-            "ln.context.uid isn't"
-            f" set in {file_path}\nCall ln.context.track() and copy/paste the output"
+            f"Cannot infer transform uid of {file_path}"
+            "\nCall `ln.track()` and copy/paste the output"
             " into the notebook"
         )
     return uid, stem_uid, version
@@ -93,7 +97,7 @@ def save_from_filepath_cli(
             if transform is None:
                 logger.error(
                     f"Did not find uid '{uid}'"
-                    " in Transform registry. Did you run ln.context.track()?"
+                    " in Transform registry. Did you run `ln.track()`?"
                 )
                 return "not-tracked-in-transform-registry"
         else:

tests/notebooks/with-title-and-initialized-consecutive.ipynb

@@ -66,8 +66,7 @@
     }
    ],
    "source": [
-    "ln.context.uid = \"hlsFXswrJjtt0000\"\n",
-    "ln.context.track()"
+    "ln.track(\"hlsFXswrJjtt0000\")"
    ]
   },
   {

tests/notebooks/with-title-and-initialized-non-consecutive.ipynb

@@ -67,7 +67,7 @@
       "source": [
         "ln.settings.transform.stem_uid = \"HDMGkxN9rgFA\"\n",
         "ln.settings.transform.version = \"1\"\n",
-        "ln.context.track()"
+        "ln.track()"
       ]
     }
   ],

tests/scripts/run-track-and-finish-sync-git.py

@@ -1,15 +1,15 @@
 import lamindb as ln
 
 ln.settings.sync_git_repo = "https://github.com/laminlabs/lamin-cli"
-ln.context.uid = "m5uCHTTpJnjQ0000"
 ln.context.name = "My good script"
+ln.track("m5uCHTTpJnjQ0000")
 
 
 if __name__ == "__main__":
     # we're using new_run here to mock the notebook situation
     # and cover the look up of an existing run in the tests
     # new_run = True is trivial
-    ln.context.track(new_run=False)
+    ln.track(new_run=False)
 
     print("hello!")
 

tests/scripts/run-track-and-finish.py

@@ -1,13 +1,14 @@
 import lamindb as ln
 
-ln.context.uid = "VFYCIuaw2GsX0000"
 ln.context.name = "My good script 2"
+ln.track("VFYCIuaw2GsX0000")
+
 
 if __name__ == "__main__":
     # we're using new_run here to mock the notebook situation
     # and cover the look up of an existing run in the tests
     # new_run = True is trivial
-    ln.context.track(new_run=False)
+    ln.track(new_run=False)
 
     print("hello!")
 

tests/test_save_notebooks.py

@@ -22,7 +22,7 @@ def test_save_not_initialized():
     )
     assert result.returncode == 1
     assert (
-        "Call ln.context.track() and copy/paste the output into the notebook"
+        "Call `ln.track()` and copy/paste the output into the notebook"
         in result.stderr.decode()
     )
 
@@ -70,7 +70,7 @@ def test_save_consecutive():
     transform = ln.Transform.filter(uid="hlsFXswrJjtt0000").one_or_none()
     assert transform is None
 
-    # let's try to save a notebook for which `ln.context.track()` was never run
+    # let's try to save a notebook for which `ln.track()` was never run
     result = subprocess.run(
         f"lamin save {notebook_path}",
         shell=True,
@@ -80,7 +80,7 @@ def test_save_consecutive():
     assert result.returncode == 1
     assert "Did not find uid 'hlsFXswrJjtt0000'" in result.stdout.decode()
 
-    # now, let's re-run this notebook so that ln.context.track() is actually run
+    # now, let's re-run this notebook so that `ln.track()` is actually run
     nbproject_test.execute_notebooks(notebook_path, print_outputs=True)
 
     # now, there is a transform record, but we're missing all artifacts
@@ -113,8 +113,8 @@ def test_save_consecutive():
 import lamindb as ln
 
 # %%
-ln.context.uid = "hlsFXswrJjtt0000"
-ln.context.track()
+ln.track("hlsFXswrJjtt0000")
+
 
 # %%
 print("my consecutive cell")

tests/test_save_scripts.py

@@ -21,7 +21,7 @@ def test_run_save_cache():
     )
     # print(result.stdout.decode())
     assert result.returncode == 1
-    assert "Did you run ln.context.track()?" in result.stdout.decode()
+    assert "Did you run ?" in result.stdout.decode()
 
     # run the script
     result = subprocess.run(