From 1edffc8354f3438aa0b2c935cf28f8a852927a54 Mon Sep 17 00:00:00 2001
From: Trevor Manz <trevor.j.manz@gmail.com>
Date: Tue, 14 Jan 2025 18:56:18 -0500
Subject: [PATCH] Add `tree` command to list notebook dependencies

---
 src/juv/__init__.py | 12 +++++++++++
 src/juv/_add.py     |  2 +-
 src/juv/_lock.py    |  2 +-
 src/juv/_remove.py  |  2 +-
 src/juv/_tree.py    | 50 +++++++++++++++++++++++++++++++++++++++++++++
 src/juv/_uv.py      | 25 +++++++++++++++++++++++
 6 files changed, 90 insertions(+), 3 deletions(-)
 create mode 100644 src/juv/_tree.py

diff --git a/src/juv/__init__.py b/src/juv/__init__.py
index 71b5524..1431784 100644
--- a/src/juv/__init__.py
+++ b/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)
diff --git a/src/juv/_add.py b/src/juv/_add.py
index de752e1..8dd0896 100644
--- a/src/juv/_add.py
+++ b/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"))
diff --git a/src/juv/_lock.py b/src/juv/_lock.py
index c10d7de..17312a8 100644
--- a/src/juv/_lock.py
+++ b/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)
 
diff --git a/src/juv/_remove.py b/src/juv/_remove.py
index 9b80893..f80005b 100644
--- a/src/juv/_remove.py
+++ b/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"))
diff --git a/src/juv/_tree.py b/src/juv/_tree.py
new file mode 100644
index 0000000..96dd368
--- /dev/null
+++ b/src/juv/_tree.py
@@ -0,0 +1,50 @@
+import tempfile
+from pathlib import Path
+
+import jupytext
+
+from ._nbutils import code_cell
+from ._pep723 import includes_inline_metadata
+from ._utils import find
+from ._uv import uv_piped
+
+
+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)
+
+        uv_piped(["tree", "--script", f.name])
+
+        if lockfile.exists():
+            notebook.metadata["uv.lock"] = lockfile.read_text(encoding="utf-8")
+            lockfile.unlink(missing_ok=True)
diff --git a/src/juv/_uv.py b/src/juv/_uv.py
index a907f40..863f982 100644
--- a/src/juv/_uv.py
+++ b/src/juv/_uv.py
@@ -2,6 +2,7 @@
 
 import os
 import subprocess
+import sys
 
 from uv import find_uv_bin
 
@@ -25,3 +26,27 @@ def uv(args: list[str], *, check: bool) -> subprocess.CompletedProcess:
     """
     uv = os.fsdecode(find_uv_bin())
     return subprocess.run([uv, *args], capture_output=True, check=check, env=os.environ)  # noqa: S603
+
+
+def uv_piped(args: list[str]) -> subprocess.CompletedProcess:
+    """Invoke a uv subprocess and pipe the result to stdout/stderr.
+
+    Parameters
+    ----------
+    args : list[str]
+        The arguments to pass to the subprocess.
+
+    Returns
+    -------
+    subprocess.CompletedProcess
+        The result of the subprocess.
+
+    """
+    uv = os.fsdecode(find_uv_bin())
+    return subprocess.run(  # noqa: S603
+        [uv, *args],
+        stdout=sys.stdout,
+        stderr=sys.stderr,
+        check=False,
+        text=True,
+    )