Notebook executor

#!/usr/bin/env python3
"""Execute the repository's notebooks without requiring Jupyter.

The runner supports the intentionally portable notebooks in this project: code
cells are ordinary Python, execute in one shared namespace, and emit stdout.
Use ``--write`` to store execution counts and stream outputs in the `.ipynb`.
"""

import argparse
import contextlib
import io
import json
import os
import sys
import traceback
from pathlib import Path
from typing import Any, Dict, List, Sequence


ROOT = Path(__file__).resolve().parents[1]
DEFAULT_NOTEBOOKS = sorted((ROOT / "notebooks").glob("*.ipynb"))


def execute_notebook(path: Path, write: bool = False) -> int:
    notebook = json.loads(path.read_text(encoding="utf-8"))
    namespace: Dict[str, Any] = {
        "__name__": "__notebook__",
        "__file__": str(path),
    }
    execution_count = 0
    original_cwd = Path.cwd()
    os.chdir(ROOT)
    try:
        for cell_index, cell in enumerate(notebook.get("cells", [])):
            if cell.get("cell_type") != "code":
                continue
            execution_count += 1
            source = cell.get("source", "")
            if isinstance(source, list):
                source = "".join(source)
            stream = io.StringIO()
            outputs: List[Dict[str, Any]] = []
            try:
                with contextlib.redirect_stdout(stream), contextlib.redirect_stderr(stream):
                    exec(compile(source, f"{path.name}:cell-{cell_index + 1}", "exec"), namespace)
            except Exception as error:  # pragma: no cover - exercised on notebook failure
                captured = stream.getvalue()
                if captured:
                    outputs.append(
                        {"name": "stdout", "output_type": "stream", "text": captured.splitlines(True)}
                    )
                outputs.append(
                    {
                        "ename": type(error).__name__,
                        "evalue": str(error),
                        "output_type": "error",
                        "traceback": traceback.format_exc().splitlines(),
                    }
                )
                cell["execution_count"] = execution_count
                cell["outputs"] = outputs
                if write:
                    path.write_text(
                        json.dumps(notebook, ensure_ascii=False, indent=1) + "\n",
                        encoding="utf-8",
                    )
                raise RuntimeError(
                    f"{path.name}: code cell {cell_index + 1} failed: {error}"
                ) from error
            captured = stream.getvalue()
            if captured:
                outputs.append(
                    {"name": "stdout", "output_type": "stream", "text": captured.splitlines(True)}
                )
            cell["execution_count"] = execution_count
            cell["outputs"] = outputs
    finally:
        os.chdir(original_cwd)
    if write:
        path.write_text(
            json.dumps(notebook, ensure_ascii=False, indent=1) + "\n",
            encoding="utf-8",
        )
    return execution_count


def parse_args(argv: Sequence[str]) -> argparse.Namespace:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("paths", nargs="*", type=Path, help="notebooks; defaults to notebooks/*.ipynb")
    parser.add_argument("--write", action="store_true", help="save execution counts and stdout")
    return parser.parse_args(argv)


def main(argv: Sequence[str] = ()) -> int:
    args = parse_args(argv or sys.argv[1:])
    paths = [path.resolve() for path in args.paths] or DEFAULT_NOTEBOOKS
    if not paths:
        print("No notebooks found", file=sys.stderr)
        return 1
    for path in paths:
        cells = execute_notebook(path, write=args.write)
        print(f"PASS {path.relative_to(ROOT)} ({cells} code cells)")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())