💚 Iteration 1

laminlabs · Sep 27, 2024 · fcba3c1 · fcba3c1
1 parent 63637e7
commit fcba3c1
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 3 deletions.
diff --git a/lamin_cli/_save.py b/lamin_cli/_save.py
@@ -21,6 +21,9 @@ def get_stem_uid_and_version_from_file(
         )
         version_pattern = re.compile(r'\.transform\.version\s*=\s*["\']([^"\']+)["\']')
     elif file_path.suffix == ".ipynb":
+        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*\\["\']([^"\']+)\\["\']'

diff --git a/tests/test_save_scripts.py b/tests/test_save_scripts.py
@@ -21,16 +21,16 @@ def test_run_save_cache():
     )
     # print(result.stdout.decode())
     assert result.returncode == 1
-    assert "Did you run ?" in result.stdout.decode()
+    assert "Did you run `ln.track()`?" in result.stdout.decode()
 
     # run the script
     result = subprocess.run(
         f"python {filepath}",
         shell=True,
         capture_output=True,
     )
-    # print(result.stdout.decode())
-    # print(result.stderr.decode())
+    print(result.stdout.decode())
+    print(result.stderr.decode())
     assert result.returncode == 0
     assert "created Transform" in result.stdout.decode()
     assert "m5uCHTTpJnjQ0000" in result.stdout.decode()