Files
mcm-26-b/scripts/run_notebook_cells.py
2026-09-13 11:45:01 +08:00

89 lines
3.4 KiB
Python

"""执行 notebook 的全部代码单元(按顺序、共享全局命名空间),用于零依赖自检。
用法::
python scripts/run_notebook_cells.py [notebook.ipynb] [--record]
等价于在 Jupyter 中"全部运行",但不需要 nbformat / nbclient:直接读 ipynb 的
JSON、拼接 code 单元、写入临时文件后交给 `python` 执行。退出码即执行结果;
``--record`` 额外把逐单元的标准输出写回 notebook 的输出字段(生成已执行副本)。
"""
from __future__ import annotations
import argparse
import contextlib
import io
import json
import subprocess
import sys
import tempfile
from pathlib import Path
ROOT = Path(__file__).resolve().parent.parent
def collect(notebook_path: Path) -> tuple[dict, list[str]]:
notebook = json.loads(notebook_path.read_text(encoding="utf-8"))
codes = ["".join(cell["source"])
for cell in notebook["cells"] if cell["cell_type"] == "code"]
return notebook, codes
def run_and_record(notebook: dict) -> None:
"""在本进程内逐单元执行并记录标准输出(写入 notebook 的 outputs 字段)。"""
scope: dict = {"__name__": "__main__"}
index = 0
for cell in notebook["cells"]:
if cell["cell_type"] != "code":
continue
buffer = io.StringIO()
with contextlib.redirect_stdout(buffer):
# 与 Jupyter 同语义:单元之间共享同一个全局命名空间
exec(compile("".join(cell["source"]), # pylint: disable=exec-used
f"<cell {index}>", "exec"), scope)
text = buffer.getvalue()
cell["execution_count"] = index + 1
cell["outputs"] = ([{"output_type": "stream", "name": "stdout",
"text": text.splitlines(keepends=True)}] if text else [])
index += 1
def main(argv: list[str] | None = None) -> int:
parser = argparse.ArgumentParser(description="执行 notebook 代码单元并自检")
parser.add_argument("notebook", nargs="?", type=Path,
default=ROOT / "mcm26b-核心代码.ipynb")
parser.add_argument("--record", action="store_true",
help="在当前进程内执行并把输出写回 notebook(生成已执行副本)")
args = parser.parse_args(argv)
notebook, codes = collect(args.notebook)
# 逐单元做语法检查,定位到具体单元
for index, cell_source in enumerate(codes):
try:
compile(cell_source, f"<cell {index}>", "exec")
except SyntaxError as exc:
print(f"单元 {index} 语法错误:{exc}", file=sys.stderr)
return 1
total_lines = sum(len(cell.splitlines()) for cell in codes)
print(f"语法检查通过:{len(codes)} 个代码单元,{total_lines} 行")
if args.record:
run_and_record(notebook)
args.notebook.write_text(
json.dumps(notebook, ensure_ascii=False, indent=1) + "\n",
encoding="utf-8")
print(f"已把逐单元输出写回 {args.notebook}")
return 0
with tempfile.TemporaryDirectory() as tmp:
target = Path(tmp) / "notebook_flat.py"
target.write_text("\n\n".join(codes), encoding="utf-8")
process = subprocess.run([sys.executable, str(target)], cwd=ROOT,
check=False)
return process.returncode
if __name__ == "__main__":
raise SystemExit(main())