File size: 2,726 Bytes
bcc8d35 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 | """Execute a percent-format source (# %% / # %% [markdown]) and write an executed .ipynb with outputs.
usage: python _m1_run.py src.py out.ipynb
Runs all code cells in ONE namespace (like a kernel), captures stdout/stderr per cell,
stores them as stream outputs. Needs nothing but the stdlib (+ whatever the notebook imports).
"""
import sys, io, json, traceback, contextlib, os, time
os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
src, out = sys.argv[1], sys.argv[2]
sys.path.insert(0, os.path.dirname(os.path.abspath(src)) + "/..")
lines = open(src).read().split("\n")
cells, cur, kind = [], [], None
def flush():
global cur, kind
if kind is not None:
body = cur[:]
while body and body[-1] == "": body.pop()
cells.append((kind, body))
cur, kind = [], None
for ln in lines:
if ln.startswith("# %%"):
flush(); kind = "markdown" if "[markdown]" in ln else "code"
else:
cur.append(ln)
flush()
ns = {"__name__": "__main__"}
nb_cells, count = [], 0
t0 = time.time()
for kind, body in cells:
if kind == "markdown":
text = "\n".join(l[2:] if l.startswith("# ") else l.lstrip("#") for l in body)
nb_cells.append({"cell_type": "markdown", "metadata": {}, "source": text.splitlines(True)})
continue
count += 1
code = "\n".join(body)
so, se = io.StringIO(), io.StringIO()
err = None
with contextlib.redirect_stdout(so), contextlib.redirect_stderr(se):
try:
exec(compile(code, f"<cell {count}>", "exec"), ns)
except BaseException:
err = traceback.format_exc()
outputs = []
if so.getvalue(): outputs.append({"output_type": "stream", "name": "stdout", "text": so.getvalue().splitlines(True)})
if se.getvalue(): outputs.append({"output_type": "stream", "name": "stderr", "text": se.getvalue().splitlines(True)})
if err:
outputs.append({"output_type": "error", "ename": "Error", "evalue": err.splitlines()[-1], "traceback": err.splitlines()})
nb_cells.append({"cell_type": "code", "execution_count": count, "metadata": {}, "source": code.splitlines(True), "outputs": outputs})
if err:
print(err, file=sys.stderr)
nb = {"cells": nb_cells, "metadata": {}, "nbformat": 4, "nbformat_minor": 4}
json.dump(nb, open(out + ".FAILED", "w"), indent=1)
sys.exit(1)
nb = {"cells": nb_cells,
"metadata": {"kernelspec": {"display_name": "Python 3", "language": "python", "name": "python3"},
"language_info": {"name": "python", "version": sys.version.split()[0]}},
"nbformat": 4, "nbformat_minor": 4}
json.dump(nb, open(out, "w"), indent=1)
print(f"wrote {out}: {count} code cells in {time.time()-t0:.1f}s")
|