ComputerAssignments/CA6_systematic_generalization/execute_notebook.py

#!/usr/bin/env python3
"""
Execute CA6 notebook with proper error handling and output capture
"""

import nbformat
from nbconvert.preprocessors import ExecutePreprocessor
import sys
import os


def execute_notebook(notebook_path, output_path=None, timeout=600):
    """
    Execute a Jupyter notebook with error handling

    Args:
        notebook_path: Path to input notebook
        output_path: Path to save executed notebook (if None, overwrites input)
        timeout: Maximum time per cell in seconds
    """
    print(f"Loading notebook: {notebook_path}")

    # Load the notebook
    with open(notebook_path, "r", encoding="utf-8") as f:
        nb = nbformat.read(f, as_version=4)

    # Create executor with increased timeout
    ep = ExecutePreprocessor(
        timeout=timeout,
        kernel_name="python3",
        allow_errors=True,  # Continue on errors
        store_widget_state=True,
    )

    print(f"Executing notebook (timeout: {timeout}s per cell)...")
    print("=" * 80)

    try:
        # Execute the notebook
        ep.preprocess(nb, {"metadata": {"path": os.path.dirname(notebook_path)}})
        print("\n" + "=" * 80)
        print("✅ Notebook execution completed!")

        # Check for any errors
        error_cells = []
        for i, cell in enumerate(nb.cells):
            if cell.cell_type == "code" and cell.get("outputs"):
                for output in cell.outputs:
                    if output.output_type == "error":
                        error_cells.append((i, cell, output))

        if error_cells:
            print(f"\n⚠️  Found {len(error_cells)} cells with errors:")
            for cell_idx, cell, error in error_cells:
                print(f"  Cell {cell_idx}: {error.get('ename', 'Unknown error')}")
                print(f"    {error.get('evalue', '')}")
        else:
            print("✅ No errors found!")

    except Exception as e:
        print(f"\n❌ Error during execution: {e}")
        import traceback

        traceback.print_exc()
        return False

    # Save the executed notebook
    if output_path is None:
        output_path = notebook_path

    print(f"\nSaving executed notebook to: {output_path}")
    with open(output_path, "w", encoding="utf-8") as f:
        nbformat.write(nb, f)

    print("✅ Notebook saved successfully!")
    return True


if __name__ == "__main__":
    notebook_file = "notebooks/CA6.ipynb"

    # Set paths
    script_dir = os.path.dirname(os.path.abspath(__file__))
    notebook_path = os.path.join(script_dir, notebook_file)
    output_path = notebook_path  # Overwrite the original

    # Execute
    success = execute_notebook(
        notebook_path, output_path, timeout=1200
    )  # 20 min per cell

    sys.exit(0 if success else 1)