analyzer/notebook_parser.py

"""
notebook_parser.py
-------------------
Extracts code cells from a .ipynb file and feeds the combined source
through the same safe AST analyzer used for .py files. No cell is ever
executed — we only read the "source" field of each code cell.
"""
import json

from .code_parser import analyze_python_source


def _strip_magics_and_shell(src):
    """Comment out Jupyter magics (%..., %%...) and shell escapes (!...)
    so the combined source is valid, parseable Python."""
    out = []
    for line in src.split("\n"):
        stripped = line.lstrip()
        if stripped.startswith("%") or stripped.startswith("!"):
            out.append("# " + line)
        else:
            out.append(line)
    return "\n".join(out)


def analyze_notebook_source(nb_text, filename="uploaded.ipynb"):
    try:
        nb = json.loads(nb_text)
    except Exception as e:
        return {"error": f"Invalid notebook JSON: {e}", "filename": filename}

    cells = nb.get("cells", [])
    chunks = []
    for cell in cells:
        if cell.get("cell_type") != "code":
            continue
        src = cell.get("source", "")
        if isinstance(src, list):
            src = "".join(src)
        chunks.append(_strip_magics_and_shell(src))

    combined = "\n\n".join(chunks)
    if not combined.strip():
        return {"error": "Notebook has no code cells", "filename": filename}

    return analyze_python_source(combined, filename=filename)