| """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") |
|
|
