Skip to content

Commit

Permalink
Add tree command to list notebook dependencies
Browse files Browse the repository at this point in the history
  • Loading branch information
manzt committed Jan 14, 2025
1 parent d65bd44 commit 1edffc8
Show file tree
Hide file tree
Showing 6 changed files with 90 additions and 3 deletions.
12 changes: 12 additions & 0 deletions src/juv/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion src/juv/_add.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
Expand Down
2 changes: 1 addition & 1 deletion src/juv/_lock.py
Original file line number Diff line number Diff line change
Expand Up @@ -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)

Expand Down
2 changes: 1 addition & 1 deletion src/juv/_remove.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"))
50 changes: 50 additions & 0 deletions src/juv/_tree.py
Original file line number Diff line number Diff line change
@@ -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)
25 changes: 25 additions & 0 deletions src/juv/_uv.py
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

import os
import subprocess
import sys

from uv import find_uv_bin

Expand All @@ -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,
)

0 comments on commit 1edffc8

Please sign in to comment.