learning/deep-probes/ast_transform.py

from __future__ import annotations

import ast
import hashlib
import json
import platform
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parents[2]
sys.path.insert(0, str(ROOT / "guest"))
from plm import transform  # type: ignore[reportMissingImports]  # noqa: E402


CASES = {
    "independent": (
        'a = tool("lookup", key="book")\n'
        'b = tool("lookup", key="shipping")\n'
        'result = a + b\n',
        {},
    ),
    "dependency": (
        'key = tool("lookup", key="choose")\n'
        'price = tool("lookup", key=key)\n'
        'shipping = tool("lookup", key="shipping")\n'
        'result = price + shipping\n',
        {},
    ),
    "branch_false": (
        'if inputs["yes"]:\n'
        '    a = tool("lookup", key="book")\n'
        'else:\n'
        '    a = tool("lookup", key="shipping")\n'
        'result = a\n',
        {"yes": False},
    ),
    "exception": (
        'try:\n'
        '    a = tool("lookup", key="book")\n'
        '    b = tool("lookup", key=inputs["absent"])\n'
        'except KeyError:\n'
        '    result = a\n',
        {},
    ),
    "unsupported_loop": (
        'result = 0\n'
        'for i in range(2):\n'
        '    result += tool("lookup", key="book")\n',
        {},
    ),
    "freshname": (
        '_spine_prepare = "user value"\n'
        'a = tool("lookup", key="book")\n'
        'result = _spine_prepare\n',
        {},
    ),
}

ADMISSION = {
    "constant_argument": 'a = tool("lookup", key="book")\n',
    "input_subscript_argument": 'a = tool("lookup", key=inputs["item"])\n',
    "data_assignment_then_tool": 'key = 1 + 2 * 3\na = tool("lookup", key=key)\n',
    "dynamic_subscript_rejected": 'a = tool("lookup", key=inputs[inputs["index"]])\n',
    "attribute_rejected": 'obj = 1\na = tool("lookup", key=obj.real)\n',
    "call_argument_rejected": 'f = 1\na = tool("lookup", key=f())\n',
    "kwargs_unpack_rejected": 'a = tool("lookup", **inputs)\n',
    "extra_positional_rejected": 'a = tool("lookup", "extra", key="book")\n',
    "multi_target_rejected": 'a = b = tool("lookup", key="book")\n',
    "loop_rejected": 'for x in []:\n    a = tool("lookup", key="book")\n',
    "import_rejected": 'import os\na = tool("lookup", key="book")\n',
    "inputs_rebind_rejected": 'inputs = {}\na = tool("lookup", key="book")\n',
    "tool_rebind_rejected": 'tool = 1\na = tool("lookup", key="book")\n',
}

VALUES = {"book": 21, "shipping": 5, "choose": "book"}


def execute_transformed(source: str, inputs: dict[str, object]):
    events: list[tuple[object, ...]] = []
    values = dict(VALUES)

    def tool(name: str, **args):
        key = args["key"]
        events.append(("call", key))
        return values[key]

    def prepare(thunk):
        try:
            name, args = thunk()
        except Exception as error:
            # The production helper returns 0 without recording this diagnostic.
            events.append(("prepare-error", type(error).__name__))
            return 0
        events.append(("prepare", args["key"]))
        return name, args

    def resolve(token, name: str, **args):
        if token == 0:
            return tool(name, **args)
        events.append(("resolve", args["key"]))
        return values[args["key"]]

    before_tree = ast.parse(source, filename="<deep-ast-before>")
    tree, helpers = transform(source)
    scope = {"__name__": "__main__", "inputs": inputs, "tool": tool}
    if helpers:
        scope.update(zip(helpers, (prepare, resolve)))
    try:
        exec(compile(tree, "<deep-ast-after>", "exec"), scope)
        result = scope.get("result")
        error = None
    except BaseException as exc:  # record, do not hide, an unexpected case failure
        result = None
        error = f"{type(exc).__name__}: {exc}"
    return {
        "before_source": ast.unparse(before_tree),
        "after_source": ast.unparse(tree),
        "before_dump": ast.dump(before_tree, indent=2),
        "after_dump": ast.dump(tree, indent=2),
        "helpers": list(helpers) if helpers else None,
        "changed": helpers is not None,
        "result": result,
        "error": error,
        "events": events,
        "dump_equal": ast.dump(tree) == ast.dump(before_tree),
    }


def write_case(name: str, source: str, inputs: dict[str, object]) -> dict[str, object]:
    out = execute_transformed(source, inputs)
    case_dir = ROOT / ".local" / "deep" / "ast" / name
    case_dir.mkdir(parents=True, exist_ok=True)
    (case_dir / "before.py").write_text(source, encoding="utf-8")
    (case_dir / "after.py").write_text(out["after_source"] + "\n", encoding="utf-8")
    (case_dir / "before.ast.txt").write_text(out["before_dump"] + "\n", encoding="utf-8")
    (case_dir / "after.ast.txt").write_text(out["after_dump"] + "\n", encoding="utf-8")
    (case_dir / "run.json").write_text(
        json.dumps(
            {
                "helpers": out["helpers"],
                "changed": out["changed"],
                "result": out["result"],
                "error": out["error"],
                "events": out["events"],
                "dump_equal": out["dump_equal"],
                "inputs": inputs,
            },
            ensure_ascii=False,
            indent=2,
        )
        + "\n",
        encoding="utf-8",
    )
    return {
        "name": name,
        "inputs": inputs,
        "before": out["before_source"],
        "after": out["after_source"],
        "helpers": out["helpers"],
        "changed": out["changed"],
        "result": out["result"],
        "error": out["error"],
        "events": out["events"],
        "dump_equal": out["dump_equal"],
        "before_ast": str((case_dir / "before.ast.txt").relative_to(ROOT)),
        "after_ast": str((case_dir / "after.ast.txt").relative_to(ROOT)),
    }


def inspect_admission(name: str, source: str) -> dict[str, object]:
    before = ast.parse(source, filename=f"<admission:{name}:before>")
    after, helpers = transform(source)
    return {
        "name": name,
        "source": source,
        "changed": helpers is not None,
        "helpers": list(helpers) if helpers else None,
        "dump_equal": ast.dump(before) == ast.dump(after),
        "after_source": ast.unparse(after),
    }


if __name__ == "__main__":
    cases = [write_case(name, source, inputs) for name, (source, inputs) in CASES.items()]
    expected = {"independent":26,"dependency":26,"branch_false":5,"exception":21,"unsupported_loop":42,"freshname":"user value"}
    assert all(case["error"] is None and case["result"] == expected[case["name"]] for case in cases), cases
    admission = [inspect_admission(name, source) for name, source in ADMISSION.items()]
    admitted = {"constant_argument", "input_subscript_argument", "data_assignment_then_tool"}
    assert all(item["changed"] == (item["name"] in admitted) for item in admission), admission
    admission_path = ROOT / ".local" / "deep" / "ast" / "admission.json"
    admission_path.write_text(json.dumps(admission, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    manifest = {
        "repository": ".",
        "reviewed_revision": "190ea2e238bba454cdcc335e95736c3f27f6be9a",
        "source_sha256": hashlib.sha256((ROOT / "guest" / "plm.py").read_bytes()).hexdigest(),
        "python": sys.version,
        "platform": platform.platform(),
        "module": "guest/plm.py",
        "cases": cases,
        "admission": str(admission_path.relative_to(ROOT)),
    }
    target = ROOT / ".local" / "deep" / "ast" / "manifest.json"
    target.write_text(json.dumps(manifest, ensure_ascii=False, indent=2) + "\n", encoding="utf-8")
    for case in cases:
        print(f"[{case['name']}] changed={case['changed']} helpers={case['helpers']} result={case['result']!r} error={case['error']!r}")
        print(f"  events={case['events']}")
        print(f"  dump_equal={case['dump_equal']}")
    for item in admission:
        print(f"[admission:{item['name']}] changed={item['changed']} helpers={item['helpers']} dump_equal={item['dump_equal']}")