nb-lab-runtime 0.1.0__py3-none-any.whl

nb_lab/__init__.py

@@ -0,0 +1,28 @@
+"""nb-lab core package."""
+
+from .cell_model import Cell, ParsedCell, assign_ids, load_sidecar, parse_cells, save_sidecar
+from .executor import PythonExecutor
+from .models import ExecutionResult, OutputItem
+from .notebook_state import NotebookState
+from .plot_meta import PlotMeta, capture_plot
+from .server import create_app
+from .watcher import PollingWatcher, WatchdogWatcher, build_watcher
+
+__all__ = [
+    "Cell",
+    "ParsedCell",
+    "assign_ids",
+    "ExecutionResult",
+    "load_sidecar",
+    "OutputItem",
+    "parse_cells",
+    "PythonExecutor",
+    "save_sidecar",
+    "PlotMeta",
+    "capture_plot",
+    "create_app",
+    "NotebookState",
+    "PollingWatcher",
+    "WatchdogWatcher",
+    "build_watcher",
+]

nb_lab/cell_model.py

@@ -0,0 +1,299 @@
+from __future__ import annotations
+
+import json
+import os
+import re
+import uuid
+from dataclasses import dataclass, field
+from difflib import SequenceMatcher
+from pathlib import Path
+from typing import Any, Literal
+
+
+CellKind = Literal["code", "markdown"]
+
+_MARKER_RE = re.compile(r"^\s*#\s*%%(?:\s*\[(?P<header>.+?)\])?\s*$")
+
+
+@dataclass
+class Cell:
+    id: str
+    kind: CellKind
+    source: str
+    start_line: int
+    end_line: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class ParsedCell:
+    kind: CellKind
+    source: str
+    start_line: int
+    end_line: int
+    metadata: dict[str, Any] = field(default_factory=dict)
+
+
+@dataclass
+class CellSnapshot:
+    id: str
+    kind: CellKind
+    source: str
+
+
+def parse_cells(text: str) -> list[ParsedCell]:
+    lines = text.splitlines(keepends=True)
+    if not lines:
+        return []
+
+    markers: list[tuple[int, CellKind, dict[str, Any]]] = []
+    for idx, line in enumerate(lines):
+        match = _MARKER_RE.match(line)
+        if not match:
+            continue
+        header = (match.group("header") or "").strip().lower()
+        kind: CellKind = "markdown" if header in {"markdown", "md"} else "code"
+        metadata: dict[str, Any] = {}
+        if header:
+            metadata["header"] = header
+        markers.append((idx, kind, metadata))
+
+    if not markers:
+        return [
+            ParsedCell(
+                kind="code",
+                source="".join(lines),
+                start_line=1,
+                end_line=len(lines),
+                metadata={},
+            )
+        ]
+
+    cells: list[ParsedCell] = []
+
+    def build_cell(
+        kind: CellKind,
+        metadata: dict[str, Any],
+        start_idx: int,
+        end_idx: int,
+        fallback_line: int,
+    ) -> ParsedCell:
+        source = "".join(lines[start_idx:end_idx])
+        if source:
+            start_line = start_idx + 1
+            end_line = end_idx
+        else:
+            start_line = fallback_line
+            end_line = fallback_line
+        return ParsedCell(
+            kind=kind,
+            source=source,
+            start_line=start_line,
+            end_line=end_line,
+            metadata=metadata,
+        )
+
+    # Handle any content before the first marker as an implicit code cell.
+    first_marker_idx, first_kind, first_meta = markers[0]
+    if first_marker_idx > 0:
+        cells.append(
+            build_cell(
+                kind="code",
+                metadata={},
+                start_idx=0,
+                end_idx=first_marker_idx,
+                fallback_line=1,
+            )
+        )
+
+    for marker_pos, (marker_idx, kind, metadata) in enumerate(markers):
+        start_idx = marker_idx + 1
+        if marker_pos + 1 < len(markers):
+            end_idx = markers[marker_pos + 1][0]
+        else:
+            end_idx = len(lines)
+        fallback_line = marker_idx + 1
+        cells.append(build_cell(kind, metadata, start_idx, end_idx, fallback_line))
+
+    return cells
+
+
+def assign_ids(
+    parsed_cells: list[ParsedCell],
+    previous: list[CellSnapshot],
+    similarity_threshold: float = 0.55,
+) -> list[Cell]:
+    remaining = list(enumerate(previous))
+    used_prev: set[int] = set()
+    cells: list[Cell] = []
+
+    for idx, parsed in enumerate(parsed_cells):
+        best_prev_index = None
+        best_score = 0.0
+        for prev_idx, prev in remaining:
+            if prev_idx in used_prev:
+                continue
+            if prev.kind != parsed.kind:
+                continue
+            score = _similarity_score(parsed.source, prev.source, parsed.kind, idx, prev_idx)
+            if score > best_score:
+                best_score = score
+                best_prev_index = prev_idx
+        if best_prev_index is not None and best_score >= similarity_threshold:
+            prev_cell = previous[best_prev_index]
+            used_prev.add(best_prev_index)
+            cell_id = prev_cell.id
+        else:
+            cell_id = uuid.uuid4().hex
+        cells.append(
+            Cell(
+                id=cell_id,
+                kind=parsed.kind,
+                source=parsed.source,
+                start_line=parsed.start_line,
+                end_line=parsed.end_line,
+                metadata=parsed.metadata,
+            )
+        )
+
+    return cells
+
+
+def _similarity_score(
+    new_source: str, old_source: str, kind: CellKind, new_idx: int, old_idx: int
+) -> float:
+    normalized_new = _normalize_source(new_source, kind)
+    normalized_old = _normalize_source(old_source, kind)
+    ratio = SequenceMatcher(None, normalized_new, normalized_old).ratio()
+    # Favor local positional continuity for stability.
+    position_penalty = 0.02 * abs(new_idx - old_idx)
+    return max(0.0, ratio - position_penalty)
+
+
+def _normalize_source(source: str, kind: CellKind) -> str:
+    lines = source.splitlines()
+    if kind == "code":
+        lines = [line for line in lines if not re.match(r"^\s*#", line)]
+    trimmed = [line.rstrip() for line in lines]
+    return "\n".join(trimmed).strip()
+
+
+def sidecar_path(notebook_path: str | Path) -> Path:
+    path = Path(notebook_path)
+    digest = _cache_key(path)
+    cache_root = Path(
+        os.environ.get("NB_LAB_CACHE_DIR", Path.home() / ".cache" / "nb-lab")
+    )
+    return cache_root / f"{digest}.json"
+
+
+def _cache_root() -> Path:
+    default_root = Path.home() / ".cache" / "nb-lab"
+    root = Path(os.environ.get("NB_LAB_CACHE_DIR", default_root))
+    if _is_writable(root):
+        return root
+    fallback = Path("/tmp/nb-lab-cache")
+    os.environ["NB_LAB_CACHE_DIR"] = str(fallback)
+    return fallback
+
+
+def _is_writable(path: Path) -> bool:
+    try:
+        path.mkdir(parents=True, exist_ok=True)
+        test_file = path / ".write_test"
+        test_file.write_text("ok", encoding="utf-8")
+        test_file.unlink(missing_ok=True)
+        return True
+    except Exception:
+        return False
+
+
+def _cache_key(path: Path) -> str:
+    return uuid.uuid5(uuid.NAMESPACE_URL, str(path.resolve())).hex
+
+
+def load_sidecar(path: str | Path) -> list[CellSnapshot]:
+    sidecar = sidecar_path(path)
+    if not sidecar.exists():
+        return []
+    raw = json.loads(sidecar.read_text(encoding="utf-8"))
+    snapshots: list[CellSnapshot] = []
+    for item in raw.get("cells", []):
+        snapshots.append(
+            CellSnapshot(
+                id=item["id"],
+                kind=item["kind"],
+                source=item["source"],
+            )
+        )
+    return snapshots
+
+
+def save_sidecar(path: str | Path, cells: list[Cell]) -> None:
+    path = Path(path)
+    sidecar = sidecar_path(path)
+    sidecar.parent.mkdir(parents=True, exist_ok=True)
+    payload = {
+        "cells": [
+            {"id": cell.id, "kind": cell.kind, "source": cell.source}
+            for cell in cells
+        ]
+    }
+    sidecar.write_text(json.dumps(payload, indent=2), encoding="utf-8")
+    meta = {"path": str(path.resolve())}
+    meta_path = sidecar.with_suffix(".meta.json")
+    meta_path.write_text(json.dumps(meta, indent=2), encoding="utf-8")
+
+
+def cache_info() -> dict[str, object]:
+    root = _cache_root()
+    if not root.exists():
+        return {"dir": str(root), "entries": 0, "size_bytes": 0}
+    entries = [p for p in root.glob("*.json") if not p.name.endswith(".meta.json")]
+    size = sum(p.stat().st_size for p in entries)
+    return {"dir": str(root), "entries": len(entries), "size_bytes": size}
+
+
+def cache_list() -> list[dict[str, object]]:
+    root = _cache_root()
+    results = []
+    if not root.exists():
+        return results
+    for sidecar in root.glob("*.json"):
+        if sidecar.name.endswith(".meta.json"):
+            continue
+        meta_path = sidecar.with_suffix(".meta.json")
+        meta = {}
+        if meta_path.exists():
+            try:
+                meta = json.loads(meta_path.read_text(encoding="utf-8"))
+            except Exception:
+                meta = {}
+        results.append(
+            {
+                "key": sidecar.stem,
+                "path": meta.get("path"),
+                "size_bytes": sidecar.stat().st_size,
+                "mtime": sidecar.stat().st_mtime,
+            }
+        )
+    return results
+
+
+def cache_clear(path: str | Path | None = None) -> int:
+    root = _cache_root()
+    if not root.exists():
+        return 0
+    deleted = 0
+    if path is None:
+        for item in root.glob("*.json"):
+            item.unlink(missing_ok=True)
+            deleted += 1
+        return deleted
+    target = sidecar_path(path)
+    meta = target.with_suffix(".meta.json")
+    for item in (target, meta):
+        if item.exists():
+            item.unlink()
+            deleted += 1
+    return deleted

nb_lab/cli.py

@@ -0,0 +1,310 @@
+from __future__ import annotations
+
+import argparse
+import json
+import sys
+from pathlib import Path
+from urllib import request
+
+import uvicorn
+
+from .cell_model import cache_clear, cache_info, cache_list
+from .server import create_app
+
+
+def _request_json(method: str, url: str, payload: dict | None = None) -> dict:
+    data = None
+    headers = {"Content-Type": "application/json"}
+    if payload is not None:
+        data = json.dumps(payload).encode("utf-8")
+    req = request.Request(url, data=data, headers=headers, method=method)
+    try:
+        with request.urlopen(req) as resp:
+            return json.loads(resp.read().decode("utf-8"))
+    except Exception as exc:
+        raise SystemExit(f"Request failed. Is the server running at {url}? ({exc})")
+
+
+def _latest_version(base_url: str) -> int:
+    data = _request_json("GET", f"{base_url}/cells")
+    return int(data.get("version", 0))
+
+
+def _cell_id_by_index(base_url: str, index: int) -> str:
+    data = _request_json("GET", f"{base_url}/cells")
+    cells = data.get("cells", [])
+    if index < 0 or index >= len(cells):
+        raise SystemExit(f"cell index {index} out of range")
+    return cells[index]["id"]
+
+
+def cmd_serve(args: argparse.Namespace) -> None:
+    _ensure_mpl_config_dir()
+    app = create_app(args.path)
+    try:
+        uvicorn.run(app, host=args.host, port=args.port)
+    except OSError as exc:
+        if args.port != 0:
+            print(f"Port {args.port} unavailable ({exc}). Retrying on an ephemeral port.")
+            uvicorn.run(app, host=args.host, port=0)
+        else:
+            raise
+
+
+def cmd_list(args: argparse.Namespace) -> None:
+    data = _request_json("GET", f"{args.base_url}/cells")
+    if args.compact:
+        print(json.dumps({"version": data.get("version"), "cells": data.get("cells", [])}, indent=2))
+    else:
+        print(json.dumps(data, indent=2))
+
+
+def cmd_run(args: argparse.Namespace) -> None:
+    if args.cell_id is None and args.cell_index is None:
+        raise SystemExit("Provide either cell_id or --cell-index")
+    if args.cell_index is not None:
+        cell_id = _cell_id_by_index(args.base_url, args.cell_index)
+    else:
+        cell_id = args.cell_id
+    base_version = (
+        args.base_version if args.base_version is not None else _latest_version(args.base_url)
+    )
+    if args.force:
+        base_version = _latest_version(args.base_url)
+    payload = {
+        "mode": args.mode,
+        "timeout_ms": args.timeout_ms,
+        "base_version": base_version,
+        "include_images": not args.no_images,
+    }
+    data = _request_json("POST", f"{args.base_url}/cells/{cell_id}/run", payload)
+    print(json.dumps(data, indent=2))
+
+def cmd_edit(args: argparse.Namespace) -> None:
+    if args.cell_id is None and args.cell_index is None:
+        raise SystemExit("Provide either cell_id or --cell-index")
+    if args.cell_index is not None:
+        cell_id = _cell_id_by_index(args.base_url, args.cell_index)
+    else:
+        cell_id = args.cell_id
+    base_version = (
+        args.base_version if args.base_version is not None else _latest_version(args.base_url)
+    )
+    if args.force:
+        base_version = _latest_version(args.base_url)
+    new_source = _resolve_new_source(args)
+    payload = {
+        "new_source": new_source,
+        "base_version": base_version,
+    }
+    data = _request_json("POST", f"{args.base_url}/cells/{cell_id}/edit", payload)
+    print(json.dumps(data, indent=2))
+
+
+def cmd_insert(args: argparse.Namespace) -> None:
+    base_version = args.base_version if args.base_version is not None else _latest_version(args.base_url)
+    payload = {
+        "after_id": args.after_id,
+        "source": args.source,
+        "kind": args.kind,
+        "base_version": base_version,
+    }
+    data = _request_json("POST", f"{args.base_url}/cells/insert_after", payload)
+    print(json.dumps(data, indent=2))
+
+
+def cmd_delete(args: argparse.Namespace) -> None:
+    if args.cell_id is None and args.cell_index is None:
+        raise SystemExit("Provide either cell_id or --cell-index")
+    if args.cell_index is not None:
+        cell_id = _cell_id_by_index(args.base_url, args.cell_index)
+    else:
+        cell_id = args.cell_id
+    base_version = (
+        args.base_version if args.base_version is not None else _latest_version(args.base_url)
+    )
+    if args.force:
+        base_version = _latest_version(args.base_url)
+    payload = {"base_version": base_version}
+    data = _request_json("DELETE", f"{args.base_url}/cells/{cell_id}", payload)
+    print(json.dumps(data, indent=2))
+
+
+def cmd_show(args: argparse.Namespace) -> None:
+    data = _request_json("GET", f"{args.base_url}/outputs/{args.cell_id}")
+    if args.compact:
+        print(json.dumps(_compact_outputs(data), indent=2))
+    else:
+        print(json.dumps(data, indent=2))
+
+
+def cmd_status(args: argparse.Namespace) -> None:
+    cells_data = _request_json("GET", f"{args.base_url}/cells")
+    status_data = _request_json("GET", f"{args.base_url}/status")
+    version = cells_data.get("version", 0)
+    cells = cells_data.get("cells", [])
+    dirty_ids = cells_data.get("dirty_ids", [])
+    output = {
+        "version": version,
+        "cell_count": len(cells),
+        "dirty_ids": dirty_ids,
+        "running_cell_id": status_data.get("running_cell_id"),
+        "running_since": status_data.get("running_since"),
+    }
+    print(json.dumps(output, indent=2))
+
+
+def cmd_cache(args: argparse.Namespace) -> None:
+    if args.cache_action == "info":
+        data = cache_info()
+        print(json.dumps(data, indent=2))
+        return
+    if args.cache_action == "list":
+        data = cache_list()
+        print(json.dumps(data, indent=2))
+        return
+    if args.cache_action == "clear":
+        if args.all and args.path:
+            raise SystemExit("Use either --all or --path, not both.")
+        if args.all:
+            args.yes = True
+            args.path = None
+        if args.path is None:
+            if not args.yes:
+                raise SystemExit("Pass --yes to clear the entire cache.")
+            deleted = cache_clear()
+        else:
+            deleted = cache_clear(args.path)
+        print(json.dumps({"deleted": deleted}, indent=2))
+
+
+def build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(prog="nb-lab")
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    serve = sub.add_parser("serve")
+    serve.add_argument("path", type=str)
+    serve.add_argument("--host", default="127.0.0.1")
+    serve.add_argument("--port", default=8787, type=int)
+    serve.set_defaults(func=cmd_serve)
+
+    list_cmd = sub.add_parser("list")
+    list_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    list_cmd.add_argument("--compact", action="store_true")
+    list_cmd.set_defaults(func=cmd_list)
+
+    run_cmd = sub.add_parser("run")
+    run_cmd.add_argument("cell_id", nargs="?")
+    run_cmd.add_argument("--cell-index", type=int)
+    run_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    run_cmd.add_argument("--base-version", type=int)
+    run_cmd.add_argument("--timeout-ms", type=int, default=60000)
+    run_cmd.add_argument(
+        "--mode",
+        default="incremental",
+        choices=["incremental", "run_above", "single"],
+    )
+    run_cmd.add_argument("--no-images", action="store_true")
+    run_cmd.add_argument("--force", action="store_true")
+    run_cmd.set_defaults(func=cmd_run)
+
+    show_cmd = sub.add_parser("show")
+    show_cmd.add_argument("cell_id")
+    show_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    show_cmd.add_argument("--compact", action="store_true")
+    show_cmd.set_defaults(func=cmd_show)
+
+    status_cmd = sub.add_parser("status")
+    status_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    status_cmd.set_defaults(func=cmd_status)
+
+    cache_cmd = sub.add_parser("cache")
+    cache_cmd.add_argument("cache_action", choices=["info", "list", "clear"])
+    cache_cmd.add_argument("--path", default=None)
+    cache_cmd.add_argument("--yes", action="store_true")
+    cache_cmd.add_argument("--all", action="store_true")
+    cache_cmd.set_defaults(func=cmd_cache)
+
+    edit_cmd = sub.add_parser("edit")
+    edit_cmd.add_argument("cell_id", nargs="?")
+    edit_cmd.add_argument("--cell-index", type=int)
+    edit_cmd.add_argument("new_source", nargs="?")
+    edit_cmd.add_argument("--from-stdin", action="store_true")
+    edit_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    edit_cmd.add_argument("--base-version", type=int)
+    edit_cmd.add_argument("--force", action="store_true")
+    edit_cmd.set_defaults(func=cmd_edit)
+
+    insert_cmd = sub.add_parser("insert")
+    insert_cmd.add_argument("--after-id", default=None)
+    insert_cmd.add_argument("--kind", default="code", choices=["code", "markdown"])
+    insert_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    insert_cmd.add_argument("--base-version", type=int)
+    insert_cmd.add_argument("source")
+    insert_cmd.set_defaults(func=cmd_insert)
+
+    delete_cmd = sub.add_parser("delete")
+    delete_cmd.add_argument("cell_id", nargs="?")
+    delete_cmd.add_argument("--cell-index", type=int)
+    delete_cmd.add_argument("--base-url", default="http://127.0.0.1:8787")
+    delete_cmd.add_argument("--base-version", type=int)
+    delete_cmd.add_argument("--force", action="store_true")
+    delete_cmd.set_defaults(func=cmd_delete)
+
+    return parser
+
+
+def main() -> None:
+    parser = build_parser()
+    args = parser.parse_args()
+    args.func(args)
+
+
+def _resolve_new_source(args: argparse.Namespace) -> str:
+    if args.from_stdin:
+        return sys.stdin.read()
+    if args.new_source is None:
+        raise SystemExit("Provide new_source or use --from-stdin")
+    return args.new_source
+
+
+def _ensure_mpl_config_dir() -> None:
+    import os
+    from pathlib import Path
+
+    if os.environ.get("MPLCONFIGDIR"):
+        return
+    default_dir = Path.home() / ".cache" / "matplotlib"
+    if _is_writable(default_dir):
+        os.environ["MPLCONFIGDIR"] = str(default_dir)
+        return
+    fallback = Path("/tmp/mpl-cache")
+    os.environ["MPLCONFIGDIR"] = str(fallback)
+
+
+def _is_writable(path: Path) -> bool:
+    try:
+        path.mkdir(parents=True, exist_ok=True)
+        test_file = path / ".write_test"
+        test_file.write_text("ok", encoding="utf-8")
+        test_file.unlink(missing_ok=True)
+        return True
+    except Exception:
+        return False
+
+
+def _compact_outputs(data: dict) -> dict:
+    result = data.get("result", {})
+    outputs = []
+    for item in result.get("outputs", []):
+        if item.get("type") == "image_png":
+            outputs.append({"type": "image_png", "data": "<omitted>"})
+        else:
+            outputs.append(item)
+    result["outputs"] = outputs
+    data["result"] = result
+    return data
+
+
+if __name__ == "__main__":
+    main()