Add tree command to list notebook dependencies (#68)

src/juv/__init__.py

@@ -477,6 +477,18 @@ def lock(
         sys.exit(1)
 
 
+@cli.command()
+@click.argument("file", type=click.Path(exists=True), required=True)
+def tree(
+    *,
+    file: str,
+) -> None:
+    """Display the notebook's dependency tree."""
+    from ._tree import tree
+
+    tree(path=Path(file))
+
+
 def main() -> None:
     """Run the CLI."""
     upgrade_legacy_jupyter_command(sys.argv)

src/juv/_add.py

@@ -141,7 +141,7 @@ def add_notebook(  # noqa: PLR0913
         cell["source"] = f.read().strip()
 
         if lockfile.exists():
-            notebook["metadata"]["uv.lock"] = lockfile.read_text()
+            notebook["metadata"]["uv.lock"] = lockfile.read_text(encoding="utf-8")
             lockfile.unlink(missing_ok=True)
 
     write_ipynb(notebook, path.with_suffix(".ipynb"))

src/juv/_lock.py

@@ -40,7 +40,7 @@ def lock(
 
         lock_file = Path(f"{temp_file.name}.lock")
 
-        notebook["metadata"]["uv.lock"] = lock_file.read_text()
+        notebook["metadata"]["uv.lock"] = lock_file.read_text(encoding="utf-8")
 
         lock_file.unlink(missing_ok=True)
 

src/juv/_remove.py

@@ -58,7 +58,7 @@ def remove(
         cell["source"] = f.read().strip()
 
         if lockfile.exists():
-            notebook["metadata"]["uv.lock"] = lockfile.read_text()
+            notebook["metadata"]["uv.lock"] = lockfile.read_text(encoding="utf-8")
             lockfile.unlink(missing_ok=True)
 
     write_ipynb(notebook, path.with_suffix(".ipynb"))

src/juv/_tree.py

@@ -0,0 +1,54 @@
+import sys
+import tempfile
+from pathlib import Path
+
+import jupytext
+
+from ._nbutils import code_cell, write_ipynb
+from ._pep723 import includes_inline_metadata
+from ._utils import find
+from ._uv import uv
+
+
+def tree(
+    path: Path,
+) -> None:
+    notebook = jupytext.read(path, fmt="ipynb")
+    lockfile_contents = notebook.get("metadata", {}).get("uv.lock")
+
+    # need a reference so we can modify the cell["source"]
+    cell = find(
+        lambda cell: (
+            cell["cell_type"] == "code"
+            and includes_inline_metadata("".join(cell["source"]))
+        ),
+        notebook["cells"],
+    )
+
+    if cell is None:
+        notebook["cells"].insert(0, code_cell("", hidden=True))
+        cell = notebook["cells"][0]
+
+    with tempfile.NamedTemporaryFile(
+        mode="w+",
+        delete=True,
+        suffix=".py",
+        dir=path.parent,
+        encoding="utf-8",
+    ) as f:
+        lockfile = Path(f"{f.name}.lock")
+
+        f.write(cell["source"].strip())
+        f.flush()
+
+        if lockfile_contents:
+            lockfile.write_text(lockfile_contents)
+
+        result = uv(["tree", "--script", f.name], check=True)
+
+        sys.stdout.write(result.stdout.decode("utf-8"))
+
+        if lockfile.exists():
+            notebook.metadata["uv.lock"] = lockfile.read_text(encoding="utf-8")
+            write_ipynb(notebook, path)
+            lockfile.unlink(missing_ok=True)

tests/test_juv.py

@@ -1156,3 +1156,68 @@ def test_remove_updates_lock(
 [options]
 exclude-newer = "2023-02-01T02:00:00Z"
 """)
+
+
+def test_tree(
+    tmp_path: pathlib.Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+
+    invoke(["init", "test.ipynb"])
+    invoke(["add", "test.ipynb", "rich"])
+    result = invoke(["tree", "test.ipynb"])
+    assert result.exit_code == 0
+    assert result.stdout == snapshot("""\
+rich v13.3.1
+├── markdown-it-py v2.1.0
+│   └── mdurl v0.1.2
+└── pygments v2.14.0
+""")
+
+
+def test_tree_updates_lock(
+    tmp_path: pathlib.Path,
+    monkeypatch: pytest.MonkeyPatch,
+) -> None:
+    monkeypatch.chdir(tmp_path)
+
+    invoke(["init", "test.ipynb"])
+    invoke(["lock", "test.ipynb"])
+    invoke(["add", "test.ipynb", "attrs"])
+    invoke(["tree", "test.ipynb"])
+    notebook = jupytext.read(tmp_path / "test.ipynb")
+    assert notebook.metadata["uv.lock"] == snapshot("""\
+version = 1
+requires-python = ">=3.13"
+
+[options]
+exclude-newer = "2023-02-01T02:00:00Z"
+
+[manifest]
+requirements = [{ name = "attrs", specifier = ">=22.2.0" }]
+
+[[package]]
+name = "attrs"
+version = "22.2.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/21/31/3f468da74c7de4fcf9b25591e682856389b3400b4b62f201e65f15ea3e07/attrs-22.2.0.tar.gz", hash = "sha256:c9227bfc2f01993c03f68db37d1d15c9690188323c067c641f1a35ca58185f99", size = 215900 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/fb/6e/6f83bf616d2becdf333a1640f1d463fef3150e2e926b7010cb0f81c95e88/attrs-22.2.0-py3-none-any.whl", hash = "sha256:29e95c7f6778868dbd49170f98f8818f78f3dc5e0e37c0b1f474e3561b240836", size = 60018 },
+]
+""")
+
+    notebook.cells[0] = new_code_cell("""# /// script
+# dependencies = []
+# requires-python = ">=3.8"
+# ///
+""")
+    write_ipynb(notebook, tmp_path / "test.ipynb")
+    invoke(["tree", "test.ipynb"])
+    assert jupytext.read(tmp_path / "test.ipynb").metadata["uv.lock"] == snapshot("""\
+version = 1
+requires-python = ">=3.8"
+
+[options]
+exclude-newer = "2023-02-01T02:00:00Z"
+""")