ComputerAssignments/CA8_neural_program_synthesis/execute_notebook.py

#!/usr/bin/env python3
"""
Script to execute notebook cells and capture outputs
"""
import json
import sys
import io
import traceback
from contextlib import redirect_stdout, redirect_stderr
import warnings

warnings.filterwarnings("ignore")


def execute_cell(code, global_vars):
    """Execute a code cell and capture output"""
    stdout_capture = io.StringIO()
    stderr_capture = io.StringIO()

    result = {"success": True, "output": "", "error": None}

    try:
        with redirect_stdout(stdout_capture), redirect_stderr(stderr_capture):
            exec(code, global_vars)

        result["output"] = stdout_capture.getvalue()
        if stderr_capture.getvalue():
            result["output"] += "\nSTDERR:\n" + stderr_capture.getvalue()

    except Exception as e:
        result["success"] = False
        result["error"] = f"{type(e).__name__}: {str(e)}\n{traceback.format_exc()}"
        result["output"] = stdout_capture.getvalue()

    return result


def main():
    notebook_path = "notebooks/CA8.ipynb"

    # Read notebook
    with open(notebook_path, "r") as f:
        notebook = json.load(f)

    # Global namespace for execution
    global_vars = {}

    # Execute cells
    cell_count = 0
    for i, cell in enumerate(notebook["cells"]):
        if cell["cell_type"] == "code":
            cell_count += 1
            print(f"\n{'='*60}")
            print(f"Executing Cell {cell_count} (Index {i})")
            print("=" * 60)

            # Get cell code
            code = "".join(cell["source"])

            # Show first few lines of code
            code_preview = "\n".join(code.split("\n")[:5])
            print(f"Code preview:\n{code_preview}...")

            # Execute
            result = execute_cell(code, global_vars)

            if result["success"]:
                print(f"✓ Success")
                if result["output"]:
                    print(f"Output:\n{result['output'][:500]}...")  # First 500 chars
            else:
                print(f"✗ Error")
                print(f"Error: {result['error'][:300]}...")

            # Add output to cell
            cell["outputs"] = [
                {
                    "output_type": "stream",
                    "name": "stdout",
                    "text": result["output"] if result["success"] else result["error"],
                }
            ]

            # Add execution count
            cell["execution_count"] = cell_count

    # Save updated notebook
    output_path = "notebooks/CA8_executed.ipynb"
    with open(output_path, "w") as f:
        json.dump(notebook, f, indent=1)

    print(f"\n\n{'='*60}")
    print(f"Notebook execution complete!")
    print(f"Saved to: {output_path}")
    print("=" * 60)


if __name__ == "__main__":
    main()