can1357
9e6eb98db5
- Removed the `_magic_cell_sh` function as `sh` cell magic was no longer required.
1172 lines
37 KiB
Python
1172 lines
37 KiB
Python
"""OMP Python runner — subprocess wrapper used by the coding-agent host.
|
|
|
|
NDJSON protocol over stdin/stdout. Host writes one JSON object per line;
|
|
wrapper writes typed frames back.
|
|
|
|
Host -> wrapper:
|
|
{"id": str, "code": str, "silent": bool?, "storeHistory": bool?}
|
|
{"id": str, "code": str, "silent": bool?, "storeHistory": bool?, "cwd": str?, "env": dict?}
|
|
{"type": "exit"} # graceful shutdown
|
|
|
|
Wrapper -> host:
|
|
{"type": "started", "id": ...}
|
|
{"type": "stdout", "id": ..., "data": str}
|
|
{"type": "stderr", "id": ..., "data": str}
|
|
{"type": "display", "id": ..., "bundle": {<mime>: <value>}}
|
|
{"type": "result", "id": ..., "bundle": {<mime>: <value>}}
|
|
{"type": "error", "id": ..., "ename": str, "evalue": str, "traceback": [str]}
|
|
{"type": "done", "id": ..., "status": "ok"|"error",
|
|
"executionCount": int, "cancelled": bool}
|
|
|
|
The runner is intentionally self-contained: no third-party imports, no IPython.
|
|
Magics are translated by a small line-scanner before AST parsing; rich display
|
|
falls back through `_repr_*_` methods so pandas/PIL/plotly etc. still render
|
|
when installed.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import asyncio
|
|
import ast
|
|
import contextvars
|
|
import base64
|
|
import builtins
|
|
import inspect
|
|
import io
|
|
import json
|
|
import os
|
|
import re
|
|
import runpy
|
|
import shlex
|
|
import signal
|
|
import subprocess
|
|
import sys
|
|
import threading
|
|
import time
|
|
import traceback
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Frame writer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
# Frames travel on a private dup of the original stdout. fd 1 itself is then
|
|
# repointed at a capture pipe: child processes spawned by user code without
|
|
# stdout=PIPE inherit fd 1, and their output is forwarded to the host as
|
|
# regular stdout frames by a drain thread instead of being written raw into
|
|
# the NDJSON channel (where it would be dropped as invalid JSON — or worse,
|
|
# spoof a frame). The wire protocol is unchanged: the host still reads NDJSON
|
|
# frames from the subprocess stdout.
|
|
_RAW_STDERR = sys.__stderr__
|
|
try:
|
|
_FRAME_FD = os.dup(sys.__stdout__.fileno())
|
|
_RAW_STDOUT = os.fdopen(_FRAME_FD, "w", encoding="utf-8", errors="backslashreplace")
|
|
_CAPTURE_READ_FD, _capture_write_fd = os.pipe()
|
|
os.dup2(_capture_write_fd, sys.__stdout__.fileno())
|
|
os.close(_capture_write_fd)
|
|
except (AttributeError, OSError, ValueError, io.UnsupportedOperation):
|
|
_RAW_STDOUT = sys.__stdout__
|
|
_CAPTURE_READ_FD = None
|
|
_OUT_LOCK = threading.Lock()
|
|
|
|
|
|
def _json_default(o: Any) -> Any:
|
|
try:
|
|
return repr(o)
|
|
except Exception:
|
|
return f"<unrepr {type(o).__name__}>"
|
|
|
|
|
|
def _emit(frame: dict) -> None:
|
|
"""Serialize a frame and write it to the host as a single NDJSON line."""
|
|
line = json.dumps(frame, ensure_ascii=False, default=_json_default)
|
|
with _OUT_LOCK:
|
|
_RAW_STDOUT.write(line)
|
|
_RAW_STDOUT.write("\n")
|
|
_RAW_STDOUT.flush()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# User stdout/stderr proxies
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _StreamProxy(io.TextIOBase):
|
|
"""Emit ``write()`` data as typed frames tied to the current request.
|
|
|
|
Writes are coalesced per request: a frame is emitted once the buffer holds
|
|
a complete line (everything up to the last newline goes out together) or
|
|
grows past ``_MAX_BUFFER`` bytes, so the common ``print()`` pair of
|
|
``write(text)`` + ``write("\\n")`` costs one frame instead of two. Partial
|
|
lines are bounded by ``flush()`` and the end-of-request flush.
|
|
"""
|
|
|
|
_MAX_BUFFER = 8192
|
|
|
|
def __init__(self, kind: str) -> None:
|
|
super().__init__()
|
|
self._kind = kind
|
|
self._lock = threading.Lock()
|
|
self._buffers: dict[str, str] = {}
|
|
|
|
def writable(self) -> bool: # noqa: D401 - protocol method
|
|
return True
|
|
|
|
def isatty(self) -> bool: # noqa: D401 - protocol method
|
|
return False
|
|
|
|
def write(self, data: Any) -> int: # type: ignore[override]
|
|
if not isinstance(data, str):
|
|
data = str(data)
|
|
if not data:
|
|
return 0
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
_RAW_STDERR.write(data)
|
|
_RAW_STDERR.flush()
|
|
return len(data)
|
|
emit_text = None
|
|
with self._lock:
|
|
buf = self._buffers.pop(rid, "") + data
|
|
if len(buf) >= self._MAX_BUFFER:
|
|
emit_text = buf
|
|
else:
|
|
nl = buf.rfind("\n")
|
|
if nl >= 0:
|
|
emit_text = buf[: nl + 1]
|
|
rest = buf[nl + 1 :]
|
|
if rest:
|
|
self._buffers[rid] = rest
|
|
else:
|
|
self._buffers[rid] = buf
|
|
if emit_text:
|
|
_emit({"type": self._kind, "id": rid, "data": emit_text})
|
|
return len(data)
|
|
|
|
def flush(self) -> None: # noqa: D401 - protocol method
|
|
rid = _CURRENT_RID.get()
|
|
if rid is not None:
|
|
self.flush_rid(rid)
|
|
return None
|
|
|
|
def flush_rid(self, rid: str) -> None:
|
|
"""Flush any buffered partial line for ``rid`` as its own frame."""
|
|
with self._lock:
|
|
buf = self._buffers.pop(rid, None)
|
|
if buf:
|
|
_emit({"type": self._kind, "id": rid, "data": buf})
|
|
|
|
|
|
def _flush_stream_proxies(rid: str) -> None:
|
|
"""Drain buffered proxy output for ``rid`` (called before its done frame)."""
|
|
for stream in (sys.stdout, sys.stderr):
|
|
if isinstance(stream, _StreamProxy):
|
|
stream.flush_rid(rid)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Runner state
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
class _RunnerState:
|
|
def __init__(self) -> None:
|
|
self.execution_count: int = 0
|
|
self.cancel_requested: bool = False
|
|
# User globals — kept across requests when running in session mode.
|
|
self.user_ns: dict[str, Any] = {
|
|
"__name__": "__main__",
|
|
"__doc__": None,
|
|
"__builtins__": builtins,
|
|
}
|
|
self.last_install_marker: int = 0
|
|
self.loop: asyncio.AbstractEventLoop | None = None
|
|
self.active_executions: int = 0
|
|
# Best-effort attribution target for captured fd-1 bytes (child
|
|
# processes inheriting stdout). With overlapping requests the most
|
|
# recently started one wins — strictly better than dropping the bytes.
|
|
self.capture_rid: str | None = None
|
|
|
|
|
|
_CURRENT_RID: contextvars.ContextVar[str | None] = contextvars.ContextVar("omp_current_rid", default=None)
|
|
_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS: contextvars.ContextVar[set[int] | None] = contextvars.ContextVar(
|
|
"omp_displayed_matplotlib_figure_ids",
|
|
default=None,
|
|
)
|
|
|
|
|
|
_STATE = _RunnerState()
|
|
|
|
|
|
def _drain_captured_stdout() -> None:
|
|
"""Forward bytes written to the captured fd 1 as stdout frames.
|
|
|
|
Runs on a daemon thread for the life of the process. Child processes that
|
|
inherit fd 1 (any ``subprocess`` call without ``stdout=PIPE``) land here.
|
|
"""
|
|
if _CAPTURE_READ_FD is None:
|
|
return
|
|
import codecs
|
|
|
|
decoder = codecs.getincrementaldecoder("utf-8")("replace")
|
|
while True:
|
|
try:
|
|
chunk = os.read(_CAPTURE_READ_FD, 65536)
|
|
except OSError:
|
|
return
|
|
if not chunk:
|
|
return
|
|
text = decoder.decode(chunk)
|
|
if not text:
|
|
continue
|
|
rid = _STATE.capture_rid
|
|
if rid is None:
|
|
_RAW_STDERR.write(text)
|
|
_RAW_STDERR.flush()
|
|
else:
|
|
_emit({"type": "stdout", "id": rid, "data": text})
|
|
|
|
|
|
def _start_capture_drain() -> None:
|
|
if _CAPTURE_READ_FD is None:
|
|
return
|
|
thread = threading.Thread(target=_drain_captured_stdout, name="omp-fd1-capture", daemon=True)
|
|
thread.start()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Magic source transformer
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_MAGIC_LINE_RE = re.compile(r"^(?P<indent>[ \t]*)(?P<name>[A-Za-z_][A-Za-z_0-9]*)(?:[ \t]+(?P<args>.*))?$")
|
|
_ASSIGN_LINE_RE = re.compile(
|
|
r"^(?P<indent>[ \t]*)(?P<lhs>[A-Za-z_][A-Za-z_0-9.\[\], ]*?)\s*=\s*(?P<rhs>.+)$"
|
|
)
|
|
|
|
|
|
def _fold_continuations(lines: list[str], start: int) -> tuple[str, int]:
|
|
"""Fold trailing backslash continuations starting at ``start``. Returns
|
|
``(folded_text, lines_consumed)``."""
|
|
parts: list[str] = []
|
|
i = start
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
if line.endswith("\\"):
|
|
parts.append(line[:-1])
|
|
i += 1
|
|
continue
|
|
parts.append(line)
|
|
i += 1
|
|
break
|
|
return ("".join(parts), i - start)
|
|
|
|
|
|
def _quote_arg(text: str) -> str:
|
|
"""Return a Python string literal that round-trips ``text`` exactly."""
|
|
return json.dumps(text, ensure_ascii=False)
|
|
|
|
|
|
def transform_cell(source: str) -> str:
|
|
"""Translate IPython-style magics + shell escapes into plain Python.
|
|
|
|
Rules
|
|
-----
|
|
* ``%name args`` -> ``__omp_magic("name", "args")``
|
|
* ``var = %name args`` -> ``var = __omp_magic("name", "args")``
|
|
* ``!cmd`` -> ``__omp_shell("cmd")``
|
|
* ``var = !cmd`` -> ``var = __omp_shell("cmd")``
|
|
* ``%%name args\\n<body>`` -> ``__omp_magic_cell("name", "args", "<body>")``
|
|
(cell magic must be the first non-whitespace token of a top-level line and
|
|
consumes the remainder of the cell)
|
|
|
|
Lines inside strings or comments are left alone — we operate on the raw
|
|
text before parsing, but the scanner only fires on the first token of each
|
|
physical line and never touches the body of triple-quoted strings because
|
|
those bodies are never first tokens themselves.
|
|
"""
|
|
|
|
if "%" not in source and "!" not in source:
|
|
return source
|
|
|
|
lines = source.splitlines()
|
|
out: list[str] = []
|
|
i = 0
|
|
while i < len(lines):
|
|
line = lines[i]
|
|
stripped = line.lstrip()
|
|
indent = line[: len(line) - len(stripped)]
|
|
|
|
# Cell magic — consumes from here to EOF.
|
|
if stripped.startswith("%%"):
|
|
head, _ = _split_magic_head(stripped[2:])
|
|
name, args = head
|
|
body_lines = lines[i + 1 :]
|
|
body = "\n".join(body_lines)
|
|
out.append(
|
|
f"{indent}__omp_magic_cell({_quote_arg(name)}, {_quote_arg(args)}, {_quote_arg(body)})"
|
|
)
|
|
return "\n".join(out)
|
|
|
|
# Line magic / shell at start of line.
|
|
if stripped.startswith("%") and not stripped.startswith("%%"):
|
|
folded, consumed = _fold_continuations(lines, i)
|
|
stripped_folded = folded.lstrip()
|
|
indent = folded[: len(folded) - len(stripped_folded)]
|
|
head, _ = _split_magic_head(stripped_folded[1:])
|
|
name, args = head
|
|
out.append(f"{indent}__omp_magic({_quote_arg(name)}, {_quote_arg(args)})")
|
|
i += consumed
|
|
continue
|
|
|
|
if stripped.startswith("!"):
|
|
folded, consumed = _fold_continuations(lines, i)
|
|
stripped_folded = folded.lstrip()
|
|
indent = folded[: len(folded) - len(stripped_folded)]
|
|
cmd = stripped_folded[1:].strip()
|
|
out.append(f"{indent}__omp_shell({_quote_arg(cmd)})")
|
|
i += consumed
|
|
continue
|
|
|
|
# Assignment forms: var = %magic / var = !cmd
|
|
m = _ASSIGN_LINE_RE.match(line)
|
|
if m:
|
|
rhs = m.group("rhs").strip()
|
|
if rhs.startswith("!"):
|
|
cmd = rhs[1:].strip()
|
|
out.append(f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_shell({_quote_arg(cmd)})")
|
|
i += 1
|
|
continue
|
|
if rhs.startswith("%") and not rhs.startswith("%%"):
|
|
head, _ = _split_magic_head(rhs[1:])
|
|
name, args = head
|
|
out.append(
|
|
f"{m.group('indent')}{m.group('lhs').rstrip()} = __omp_magic({_quote_arg(name)}, {_quote_arg(args)})"
|
|
)
|
|
i += 1
|
|
continue
|
|
|
|
out.append(line)
|
|
i += 1
|
|
|
|
return "\n".join(out)
|
|
|
|
|
|
def _split_magic_head(text: str) -> tuple[tuple[str, str], str]:
|
|
"""Split ``"name rest"`` into ``("name", "rest")``."""
|
|
text = text.lstrip()
|
|
if not text:
|
|
return ("", ""), ""
|
|
m = re.match(r"([A-Za-z_][A-Za-z_0-9]*)(?:\s+(.*))?$", text)
|
|
if not m:
|
|
return ("", text), ""
|
|
return (m.group(1), (m.group(2) or "").rstrip()), ""
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Magic registry
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_LINE_MAGICS: dict[str, Callable[[str], Any]] = {}
|
|
_CELL_MAGICS: dict[str, Callable[[str, str], Any]] = {}
|
|
|
|
|
|
def line_magic(name: str) -> Callable[[Callable[[str], Any]], Callable[[str], Any]]:
|
|
def decorator(fn: Callable[[str], Any]) -> Callable[[str], Any]:
|
|
_LINE_MAGICS[name] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def cell_magic(name: str) -> Callable[[Callable[[str, str], Any]], Callable[[str, str], Any]]:
|
|
def decorator(fn: Callable[[str, str], Any]) -> Callable[[str, str], Any]:
|
|
_CELL_MAGICS[name] = fn
|
|
return fn
|
|
|
|
return decorator
|
|
|
|
|
|
def _emit_status(op: str, **data: Any) -> None:
|
|
bundle = {"application/x-omp-status": {"op": op, **data}}
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
return
|
|
_emit({"type": "display", "id": rid, "bundle": bundle})
|
|
|
|
|
|
@line_magic("pip")
|
|
def _magic_pip(args: str) -> None:
|
|
argv = shlex.split(args) if args else ["--help"]
|
|
cmd = [sys.executable, "-m", "pip", *argv]
|
|
proc = subprocess.Popen(
|
|
cmd,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
installed_packages: list[str] = []
|
|
assert proc.stdout is not None
|
|
for raw_line in proc.stdout:
|
|
sys.stdout.write(raw_line)
|
|
m = re.search(r"Successfully installed\s+(.+)$", raw_line)
|
|
if m:
|
|
for token in m.group(1).split():
|
|
# Token is name-version; drop the version suffix.
|
|
pkg = token.rsplit("-", 1)[0]
|
|
installed_packages.append(pkg.replace("_", "-"))
|
|
proc.wait()
|
|
if installed_packages:
|
|
import importlib
|
|
|
|
importlib.invalidate_caches()
|
|
prefixes = {pkg.lower().replace("-", "_") for pkg in installed_packages}
|
|
for mod_name in list(sys.modules):
|
|
head = mod_name.split(".", 1)[0].lower()
|
|
if head in prefixes:
|
|
sys.modules.pop(mod_name, None)
|
|
_emit_status("pip", args=args, installed=installed_packages, exit_code=proc.returncode)
|
|
|
|
|
|
@line_magic("cd")
|
|
def _magic_cd(args: str) -> str:
|
|
path = os.path.expanduser(args.strip()) or os.path.expanduser("~")
|
|
os.chdir(path)
|
|
cwd = os.getcwd()
|
|
_emit_status("cd", path=cwd)
|
|
return cwd
|
|
|
|
|
|
@line_magic("pwd")
|
|
def _magic_pwd(_args: str) -> str:
|
|
cwd = os.getcwd()
|
|
_emit_status("pwd", path=cwd)
|
|
return cwd
|
|
|
|
|
|
@line_magic("ls")
|
|
def _magic_ls(args: str) -> list[str]:
|
|
target = os.path.expanduser(args.strip()) or "."
|
|
entries = sorted(os.listdir(target))
|
|
_emit_status("ls", path=os.path.abspath(target), count=len(entries))
|
|
return entries
|
|
|
|
|
|
@line_magic("env")
|
|
def _magic_env(args: str) -> Any:
|
|
args = args.strip()
|
|
if not args:
|
|
return dict(sorted(os.environ.items()))
|
|
if "=" in args:
|
|
key, value = args.split("=", 1)
|
|
os.environ[key.strip()] = value.strip()
|
|
return value.strip()
|
|
return os.environ.get(args)
|
|
|
|
|
|
@line_magic("set_env")
|
|
def _magic_set_env(args: str) -> str:
|
|
parts = args.split(None, 1)
|
|
if len(parts) != 2:
|
|
raise ValueError("Usage: %set_env KEY VALUE")
|
|
key, value = parts
|
|
os.environ[key] = value
|
|
return value
|
|
|
|
|
|
@line_magic("time")
|
|
def _magic_time(args: str) -> Any:
|
|
start = time.perf_counter()
|
|
result = eval(args, _STATE.user_ns)
|
|
elapsed = time.perf_counter() - start
|
|
sys.stdout.write(f"Wall time: {elapsed * 1000:.2f} ms\n")
|
|
_emit_status("time", elapsed_ms=round(elapsed * 1000, 3))
|
|
return result
|
|
|
|
|
|
@line_magic("timeit")
|
|
def _magic_timeit(args: str) -> None:
|
|
import timeit as _timeit
|
|
|
|
timer = _timeit.Timer(stmt=args, globals=_STATE.user_ns)
|
|
iters, total = timer.autorange()
|
|
per = total / iters
|
|
sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
|
|
_emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
|
|
|
|
|
|
@line_magic("who")
|
|
def _magic_who(_args: str) -> list[str]:
|
|
names = sorted(
|
|
name
|
|
for name, value in _STATE.user_ns.items()
|
|
if not name.startswith("_") and not callable(value) or hasattr(value, "__class__")
|
|
)
|
|
return [n for n in names if not n.startswith("__")]
|
|
|
|
|
|
@line_magic("whos")
|
|
def _magic_whos(_args: str) -> list[tuple[str, str]]:
|
|
rows = []
|
|
for name in sorted(_STATE.user_ns):
|
|
if name.startswith("__"):
|
|
continue
|
|
value = _STATE.user_ns[name]
|
|
rows.append((name, type(value).__name__))
|
|
return rows
|
|
|
|
|
|
@line_magic("reset")
|
|
def _magic_reset(_args: str) -> None:
|
|
_STATE.user_ns.clear()
|
|
_STATE.user_ns.update({"__name__": "__main__", "__doc__": None, "__builtins__": builtins})
|
|
_install_builtins(_STATE.user_ns)
|
|
_emit_status("reset")
|
|
|
|
|
|
@line_magic("load")
|
|
def _magic_load(args: str) -> None:
|
|
path = Path(os.path.expanduser(args.strip()))
|
|
source = path.read_text(encoding="utf-8")
|
|
_emit({"type": "display", "id": _CURRENT_RID.get(), "bundle": {"text/plain": source}})
|
|
_exec_source(source, _STATE.user_ns)
|
|
|
|
|
|
@line_magic("run")
|
|
def _magic_run(args: str) -> None:
|
|
parts = shlex.split(args) if args else []
|
|
if not parts:
|
|
raise ValueError("Usage: %run <path>")
|
|
target = os.path.expanduser(parts[0])
|
|
saved_argv = sys.argv
|
|
try:
|
|
sys.argv = [target, *parts[1:]]
|
|
result_ns = runpy.run_path(target, run_name="__main__")
|
|
finally:
|
|
sys.argv = saved_argv
|
|
for name, value in result_ns.items():
|
|
if name.startswith("__"):
|
|
continue
|
|
_STATE.user_ns[name] = value
|
|
|
|
|
|
@cell_magic("bash")
|
|
def _magic_cell_bash(args: str, body: str) -> int:
|
|
return _run_shell_body(body, shell_arg="/bin/bash")
|
|
|
|
@cell_magic("capture")
|
|
def _magic_cell_capture(args: str, body: str) -> str:
|
|
"""Capture stdout/stderr of body; bind to ``args`` (a name) if provided."""
|
|
captured = io.StringIO()
|
|
saved_stdout, saved_stderr = sys.stdout, sys.stderr
|
|
sys.stdout = sys.stderr = captured
|
|
try:
|
|
_exec_source(body, _STATE.user_ns)
|
|
finally:
|
|
sys.stdout, sys.stderr = saved_stdout, saved_stderr
|
|
text = captured.getvalue()
|
|
name = args.strip()
|
|
if name:
|
|
_STATE.user_ns[name] = text
|
|
return text
|
|
|
|
|
|
@cell_magic("timeit")
|
|
def _magic_cell_timeit(args: str, body: str) -> None:
|
|
import timeit as _timeit
|
|
|
|
timer = _timeit.Timer(stmt=body, globals=_STATE.user_ns)
|
|
iters, total = timer.autorange()
|
|
per = total / iters
|
|
sys.stdout.write(f"{iters} loops, best of 1: {per * 1e6:.2f} us per loop\n")
|
|
_emit_status("timeit", loops=iters, total_ms=round(total * 1000, 3))
|
|
|
|
|
|
@cell_magic("writefile")
|
|
def _magic_cell_writefile(args: str, body: str) -> str:
|
|
path = Path(os.path.expanduser(args.strip()))
|
|
path.parent.mkdir(parents=True, exist_ok=True)
|
|
path.write_text(body, encoding="utf-8")
|
|
_emit_status("writefile", path=str(path), bytes=len(body))
|
|
return str(path)
|
|
|
|
|
|
def _run_shell_body(body: str, *, shell_arg: str) -> int:
|
|
proc = subprocess.Popen(
|
|
[shell_arg, "-c", body],
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
bufsize=1,
|
|
)
|
|
assert proc.stdout is not None
|
|
for raw_line in proc.stdout:
|
|
sys.stdout.write(raw_line)
|
|
proc.wait()
|
|
return proc.returncode
|
|
|
|
|
|
def __omp_magic(name: str, args: str) -> Any:
|
|
fn = _LINE_MAGICS.get(name)
|
|
if fn is None:
|
|
raise NameError(f"UsageError: Line magic function '%{name}' not found.")
|
|
return fn(args)
|
|
|
|
|
|
def __omp_magic_cell(name: str, args: str, body: str) -> Any:
|
|
fn = _CELL_MAGICS.get(name)
|
|
if fn is None:
|
|
raise NameError(f"UsageError: Cell magic function '%%{name}' not found.")
|
|
return fn(args, body)
|
|
|
|
|
|
class _ShellResult(list):
|
|
"""Result of ``!cmd`` — list of stripped output lines."""
|
|
|
|
def __init__(self, lines: list[str], returncode: int) -> None:
|
|
super().__init__(lines)
|
|
self.returncode = returncode
|
|
|
|
@property
|
|
def n(self) -> str: # IPython compat
|
|
return "\n".join(self)
|
|
|
|
@property
|
|
def s(self) -> str: # IPython compat
|
|
return " ".join(self)
|
|
|
|
|
|
def __omp_shell(cmd: str) -> _ShellResult:
|
|
proc = subprocess.run(
|
|
cmd,
|
|
shell=True,
|
|
stdout=subprocess.PIPE,
|
|
stderr=subprocess.STDOUT,
|
|
text=True,
|
|
)
|
|
if proc.stdout:
|
|
sys.stdout.write(proc.stdout)
|
|
lines = [line for line in (proc.stdout or "").splitlines()]
|
|
return _ShellResult(lines, proc.returncode)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Display dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_REPR_MIMES = [
|
|
("_repr_html_", "text/html"),
|
|
("_repr_markdown_", "text/markdown"),
|
|
("_repr_svg_", "image/svg+xml"),
|
|
("_repr_png_", "image/png"),
|
|
("_repr_jpeg_", "image/jpeg"),
|
|
("_repr_json_", "application/json"),
|
|
("_repr_latex_", "text/latex"),
|
|
]
|
|
|
|
|
|
def _is_matplotlib_figure(value: Any) -> bool:
|
|
figure_module = sys.modules.get("matplotlib.figure")
|
|
figure_cls = getattr(figure_module, "Figure", None)
|
|
if isinstance(figure_cls, type) and isinstance(value, figure_cls):
|
|
return True
|
|
|
|
value_type = type(value)
|
|
return value_type.__module__ == "matplotlib.figure" and value_type.__name__ == "Figure"
|
|
|
|
|
|
def _matplotlib_figure_png(value: Any) -> str | None:
|
|
if not _is_matplotlib_figure(value):
|
|
return None
|
|
|
|
savefig = getattr(value, "savefig", None)
|
|
if not callable(savefig):
|
|
return None
|
|
|
|
try:
|
|
buf = io.BytesIO()
|
|
savefig(buf, format="png", bbox_inches="tight")
|
|
except Exception:
|
|
return None
|
|
|
|
displayed_ids = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get()
|
|
if displayed_ids is not None:
|
|
displayed_ids.add(id(value))
|
|
return base64.b64encode(buf.getvalue()).decode("ascii")
|
|
|
|
|
|
def _coerce_image_bytes(value: Any) -> str:
|
|
if isinstance(value, (bytes, bytearray)):
|
|
return base64.b64encode(bytes(value)).decode("ascii")
|
|
if isinstance(value, str):
|
|
return value
|
|
return base64.b64encode(repr(value).encode("utf-8")).decode("ascii")
|
|
|
|
|
|
def _mime_bundle(value: Any) -> dict:
|
|
"""Build a Jupyter-style MIME bundle for ``value``.
|
|
|
|
Honors ``_repr_mimebundle_`` first, falls back to individual ``_repr_*_``
|
|
accessors, and always provides ``text/plain``.
|
|
"""
|
|
bundle: dict[str, Any] = {}
|
|
matplotlib_png = _matplotlib_figure_png(value)
|
|
if matplotlib_png is not None:
|
|
bundle["image/png"] = matplotlib_png
|
|
|
|
|
|
mimebundle = getattr(value, "_repr_mimebundle_", None)
|
|
if callable(mimebundle):
|
|
try:
|
|
data = mimebundle()
|
|
except Exception:
|
|
data = None
|
|
if isinstance(data, tuple):
|
|
data = data[0]
|
|
if isinstance(data, dict):
|
|
bundle.update({str(k): v for k, v in data.items()})
|
|
|
|
for attr, mime in _REPR_MIMES:
|
|
if mime in bundle:
|
|
continue
|
|
repr_fn = getattr(value, attr, None)
|
|
if not callable(repr_fn):
|
|
continue
|
|
try:
|
|
data = repr_fn()
|
|
except Exception:
|
|
continue
|
|
if data is None:
|
|
continue
|
|
if mime in ("image/png", "image/jpeg"):
|
|
bundle[mime] = _coerce_image_bytes(data)
|
|
else:
|
|
bundle[mime] = data
|
|
|
|
if "text/plain" not in bundle:
|
|
try:
|
|
bundle["text/plain"] = repr(value)
|
|
except Exception:
|
|
bundle["text/plain"] = f"<unrepr {type(value).__name__}>"
|
|
|
|
return bundle
|
|
|
|
|
|
def _emit_display(bundle: dict, *, kind: str = "display") -> None:
|
|
rid = _CURRENT_RID.get()
|
|
if rid is None:
|
|
return
|
|
_emit({"type": kind, "id": rid, "bundle": bundle})
|
|
|
|
|
|
def __omp_display(value: Any, *, raw: bool = False, kind: str = "display") -> None:
|
|
if raw:
|
|
if not isinstance(value, dict):
|
|
raise TypeError("display(..., raw=True) requires a MIME bundle dict")
|
|
bundle = {str(k): v for k, v in value.items()}
|
|
if "text/plain" not in bundle:
|
|
bundle["text/plain"] = ""
|
|
_emit_display(bundle, kind=kind)
|
|
return
|
|
_emit_display(_mime_bundle(value), kind=kind)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Matplotlib post-cell flush
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _flush_matplotlib_figures() -> None:
|
|
plt = sys.modules.get("matplotlib.pyplot")
|
|
if plt is None:
|
|
return
|
|
try:
|
|
fignums = list(plt.get_fignums())
|
|
except Exception:
|
|
return
|
|
for num in fignums:
|
|
try:
|
|
fig = plt.figure(num)
|
|
if id(fig) in (_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.get() or set()):
|
|
plt.close(fig)
|
|
continue
|
|
buf = io.BytesIO()
|
|
fig.savefig(buf, format="png", bbox_inches="tight")
|
|
data = base64.b64encode(buf.getvalue()).decode("ascii")
|
|
_emit_display({"image/png": data, "text/plain": f"<Figure {num}>"})
|
|
plt.close(fig)
|
|
except Exception:
|
|
continue
|
|
|
|
|
|
# Force a non-interactive backend before user code imports matplotlib. Set as
|
|
# environ default so the user can still override it explicitly.
|
|
os.environ.setdefault("MPLBACKEND", "Agg")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Builtin injection
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _install_builtins(ns: dict) -> None:
|
|
ns["display"] = __omp_display
|
|
ns["__omp_display"] = __omp_display
|
|
ns["__omp_magic"] = __omp_magic
|
|
ns["__omp_magic_cell"] = __omp_magic_cell
|
|
ns["__omp_shell"] = __omp_shell
|
|
ns["__omp_current_run_id__"] = lambda: _CURRENT_RID.get()
|
|
|
|
|
|
_install_builtins(_STATE.user_ns)
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Source execution (split last expression for rich display)
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
_TLA_FLAG = getattr(ast, "PyCF_ALLOW_TOP_LEVEL_AWAIT", 0x2000)
|
|
|
|
|
|
def _await_sync(coro) -> Any:
|
|
try:
|
|
running_loop = asyncio.get_running_loop()
|
|
except RuntimeError:
|
|
running_loop = None
|
|
if running_loop is not None and running_loop.is_running():
|
|
raise RuntimeError("top-level await is not supported from synchronous magic execution")
|
|
return asyncio.run(coro)
|
|
|
|
|
|
def _run_compiled_sync(code, ns: dict, *, want_value: bool) -> Any:
|
|
"""Synchronous execution path used by nested magic helpers."""
|
|
if code.co_flags & inspect.CO_COROUTINE:
|
|
result = _await_sync(eval(code, ns))
|
|
return result if want_value else None
|
|
if want_value:
|
|
return eval(code, ns)
|
|
exec(code, ns)
|
|
return None
|
|
|
|
|
|
|
|
async def _run_compiled_async(code, ns: dict, *, want_value: bool) -> Any:
|
|
"""Execute a code object in the persistent event loop.
|
|
|
|
Coroutine code is awaited in this task so top-level ``await`` interleaves
|
|
with sibling requests. Plain statement/expression code runs on the main
|
|
runner thread so SIGINT can interrupt it reliably.
|
|
"""
|
|
if code.co_flags & inspect.CO_COROUTINE:
|
|
result = await eval(code, ns)
|
|
return result if want_value else None
|
|
if want_value:
|
|
return eval(code, ns)
|
|
exec(code, ns)
|
|
return None
|
|
|
|
|
|
def _compile_source(source: str) -> tuple[Any, Any | None, bool]:
|
|
module = ast.parse(source, mode="exec")
|
|
if not module.body:
|
|
return None, None, False
|
|
|
|
last = module.body[-1]
|
|
if isinstance(last, ast.Expr):
|
|
body_module = ast.Module(body=module.body[:-1], type_ignores=[])
|
|
expr_module = ast.Expression(body=last.value)
|
|
ast.copy_location(expr_module, last)
|
|
body_code = compile(body_module, "<cell>", "exec", flags=_TLA_FLAG)
|
|
expr_code = compile(expr_module, "<cell>", "eval", flags=_TLA_FLAG)
|
|
return body_code, expr_code, True
|
|
|
|
return compile(module, "<cell>", "exec", flags=_TLA_FLAG), None, False
|
|
|
|
|
|
def _exec_source(source: str, ns: dict) -> None:
|
|
"""Synchronous source execution for legacy magic helpers."""
|
|
body_code, expr_code, has_expr = _compile_source(source)
|
|
if body_code is None:
|
|
return
|
|
_run_compiled_sync(body_code, ns, want_value=False)
|
|
if has_expr and expr_code is not None:
|
|
value = _run_compiled_sync(expr_code, ns, want_value=True)
|
|
if value is not None:
|
|
__omp_display(value, kind="result")
|
|
|
|
|
|
async def _exec_source_async(source: str, ns: dict) -> None:
|
|
"""Compile + execute ``source``; if the last node is an expression, route
|
|
its value through ``__omp_display`` so dataframes/figures render rich.
|
|
Top-level ``await`` / ``async for`` / ``async with`` is permitted; awaited
|
|
regions yield to other requests in the runner's persistent event loop."""
|
|
body_code, expr_code, has_expr = _compile_source(source)
|
|
if body_code is None:
|
|
return
|
|
await _run_compiled_async(body_code, ns, want_value=False)
|
|
if has_expr and expr_code is not None:
|
|
value = await _run_compiled_async(expr_code, ns, want_value=True)
|
|
if value is not None:
|
|
__omp_display(value, kind="result")
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Signal handling
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _install_idle_sigint() -> None:
|
|
try:
|
|
signal.signal(signal.SIGINT, signal.SIG_IGN)
|
|
except (OSError, ValueError):
|
|
# Some platforms (Windows in non-console mode) reject this; fine.
|
|
pass
|
|
|
|
|
|
def _install_exec_sigint() -> None:
|
|
try:
|
|
signal.signal(signal.SIGINT, signal.default_int_handler)
|
|
except (OSError, ValueError):
|
|
pass
|
|
|
|
|
|
def _begin_exec_sigint() -> None:
|
|
_STATE.active_executions += 1
|
|
_install_exec_sigint()
|
|
|
|
|
|
def _end_exec_sigint() -> None:
|
|
if _STATE.active_executions > 0:
|
|
_STATE.active_executions -= 1
|
|
if _STATE.active_executions == 0:
|
|
_install_idle_sigint()
|
|
|
|
|
|
_MANAGED_ENV_KEYS = (
|
|
"PI_SESSION_FILE",
|
|
"PI_ARTIFACTS_DIR",
|
|
"PI_TOOL_BRIDGE_URL",
|
|
"PI_TOOL_BRIDGE_TOKEN",
|
|
"PI_TOOL_BRIDGE_SESSION",
|
|
"PI_EVAL_LOCAL_ROOTS",
|
|
)
|
|
|
|
|
|
def _apply_request_runtime(req: dict) -> None:
|
|
cwd = req.get("cwd")
|
|
if isinstance(cwd, str) and cwd:
|
|
os.chdir(cwd)
|
|
try:
|
|
sys.path.remove(cwd)
|
|
except ValueError:
|
|
pass
|
|
sys.path.insert(0, cwd)
|
|
|
|
env = req.get("env")
|
|
if isinstance(env, dict):
|
|
for key in _MANAGED_ENV_KEYS:
|
|
value = env.get(key)
|
|
if isinstance(value, str):
|
|
os.environ[key] = value
|
|
elif value is None:
|
|
os.environ.pop(key, None)
|
|
|
|
def _start_parent_watchdog() -> None:
|
|
"""Self-terminate when the host process dies.
|
|
|
|
The main loop only exits when stdin EOFs, which only happens once user
|
|
code finishes and the next ``readline`` call returns. If the host gets
|
|
SIGKILL mid-execution (or any way that skips graceful shutdown) the
|
|
runner would otherwise outlive its parent and keep holding kernel
|
|
state. Poll ``os.getppid()`` instead and ``os._exit`` the moment we get
|
|
reparented \u2014 covers POSIX hosts. Windows has no reliable ppid
|
|
equivalent; there we still bail out on the next stdin read.
|
|
"""
|
|
if os.name != "posix":
|
|
return
|
|
original_ppid = os.getppid()
|
|
if original_ppid <= 1:
|
|
return
|
|
|
|
def watch() -> None:
|
|
while True:
|
|
try:
|
|
if os.getppid() != original_ppid:
|
|
os._exit(0)
|
|
except Exception:
|
|
return
|
|
time.sleep(10)
|
|
|
|
thread = threading.Thread(target=watch, name="omp-parent-watchdog", daemon=True)
|
|
thread.start()
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Request dispatch
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
async def _handle_request_async(req: dict) -> None:
|
|
rid = str(req.get("id"))
|
|
token = _CURRENT_RID.set(rid)
|
|
displayed_matplotlib_token = _CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.set(set())
|
|
_STATE.capture_rid = rid
|
|
_STATE.user_ns["__omp_run_id__"] = rid
|
|
_STATE.cancel_requested = False
|
|
_STATE.execution_count += 1
|
|
execution_count = _STATE.execution_count
|
|
_emit({"type": "started", "id": rid})
|
|
|
|
status: str = "ok"
|
|
cancelled = False
|
|
|
|
try:
|
|
try:
|
|
_apply_request_runtime(req)
|
|
transformed = transform_cell(req.get("code", ""))
|
|
except SyntaxError as exc:
|
|
_emit_error(rid, exc)
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": "error",
|
|
"executionCount": execution_count,
|
|
"cancelled": False,
|
|
})
|
|
return
|
|
except BaseException as exc: # noqa: BLE001 - runtime setup errors must settle the request
|
|
_emit_error(rid, exc)
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": "error",
|
|
"executionCount": execution_count,
|
|
"cancelled": False,
|
|
})
|
|
return
|
|
|
|
_begin_exec_sigint()
|
|
try:
|
|
await _exec_source_async(transformed, _STATE.user_ns)
|
|
except KeyboardInterrupt:
|
|
cancelled = True
|
|
status = "error"
|
|
_emit_error(rid, KeyboardInterrupt("Execution interrupted"))
|
|
except SystemExit as exc:
|
|
status = "error"
|
|
_emit_error(rid, exc)
|
|
except BaseException as exc: # noqa: BLE001 - we want to surface every user error
|
|
status = "error"
|
|
_emit_error(rid, exc)
|
|
finally:
|
|
_end_exec_sigint()
|
|
try:
|
|
_flush_matplotlib_figures()
|
|
except Exception:
|
|
pass
|
|
|
|
_flush_stream_proxies(rid)
|
|
_emit({
|
|
"type": "done",
|
|
"id": rid,
|
|
"status": status,
|
|
"executionCount": execution_count,
|
|
"cancelled": cancelled,
|
|
})
|
|
finally:
|
|
if _STATE.capture_rid == rid:
|
|
_STATE.capture_rid = None
|
|
_flush_stream_proxies(rid)
|
|
_CURRENT_RID.reset(token)
|
|
_CURRENT_DISPLAYED_MATPLOTLIB_FIGURE_IDS.reset(displayed_matplotlib_token)
|
|
|
|
|
|
def _emit_error(rid: str, exc: BaseException) -> None:
|
|
tb_lines = traceback.format_exception(type(exc), exc, exc.__traceback__)
|
|
_emit({
|
|
"type": "error",
|
|
"id": rid,
|
|
"ename": type(exc).__name__,
|
|
"evalue": str(exc),
|
|
"traceback": [line.rstrip("\n") for line in tb_lines],
|
|
})
|
|
|
|
|
|
# ---------------------------------------------------------------------------
|
|
# Main loop
|
|
# ---------------------------------------------------------------------------
|
|
|
|
|
|
def _read_stdin(loop: asyncio.AbstractEventLoop, queue: asyncio.Queue, stdin) -> None:
|
|
for raw_line in stdin:
|
|
line = raw_line.strip()
|
|
if not line:
|
|
continue
|
|
try:
|
|
req = json.loads(line)
|
|
except json.JSONDecodeError as exc:
|
|
_emit({
|
|
"type": "error",
|
|
"id": "",
|
|
"ename": "ProtocolError",
|
|
"evalue": f"Invalid JSON request: {exc}",
|
|
"traceback": [],
|
|
})
|
|
continue
|
|
loop.call_soon_threadsafe(queue.put_nowait, req)
|
|
loop.call_soon_threadsafe(queue.put_nowait, {"type": "exit"})
|
|
|
|
|
|
async def _main_async() -> None:
|
|
sys.stdout = _StreamProxy("stdout")
|
|
sys.stderr = _StreamProxy("stderr")
|
|
_install_idle_sigint()
|
|
_start_parent_watchdog()
|
|
_start_capture_drain()
|
|
|
|
stdin = sys.__stdin__
|
|
if stdin is None:
|
|
return
|
|
|
|
loop = asyncio.get_running_loop()
|
|
_STATE.loop = loop
|
|
queue: asyncio.Queue = asyncio.Queue()
|
|
reader = threading.Thread(target=_read_stdin, args=(loop, queue, stdin), name="omp-stdin-reader", daemon=True)
|
|
reader.start()
|
|
|
|
tasks: set[asyncio.Task] = set()
|
|
def _task_done(task: asyncio.Task) -> None:
|
|
tasks.discard(task)
|
|
try:
|
|
exc = task.exception()
|
|
except asyncio.CancelledError:
|
|
return
|
|
if exc is not None:
|
|
_emit_error("", exc)
|
|
try:
|
|
while True:
|
|
req = await queue.get()
|
|
if req.get("type") == "exit":
|
|
break
|
|
task = asyncio.create_task(_handle_request_async(req))
|
|
tasks.add(task)
|
|
task.add_done_callback(_task_done)
|
|
finally:
|
|
for task in tasks:
|
|
task.cancel()
|
|
if tasks:
|
|
await asyncio.gather(*tasks, return_exceptions=True)
|
|
|
|
|
|
def main() -> None:
|
|
asyncio.run(_main_async())
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|