tools/python_exec.py

"""Constrained execution for small benchmark Python attachments."""

from __future__ import annotations

import ast
import subprocess
import sys
from pathlib import Path

_ALLOWED_IMPORTS = {
    "cmath",
    "collections",
    "decimal",
    "fractions",
    "functools",
    "itertools",
    "math",
    "operator",
    "random",
    "re",
    "statistics",
    "time",
}
_BLOCKED_CALLS = {
    "breakpoint",
    "compile",
    "eval",
    "exec",
    "input",
    "open",
    "__import__",
}


def _validate(source: str) -> None:
    tree = ast.parse(source)
    for node in ast.walk(tree):
        if isinstance(node, (ast.Import, ast.ImportFrom)):
            names = (
                [a.name.split(".")[0] for a in node.names]
                if isinstance(node, ast.Import)
                else [(node.module or "").split(".")[0]]
            )
            if any(name not in _ALLOWED_IMPORTS for name in names):
                raise ValueError(
                    f"Blocked import in Python attachment: {', '.join(names)}"
                )
        if (
            isinstance(node, ast.Call)
            and isinstance(node.func, ast.Name)
            and node.func.id in _BLOCKED_CALLS
        ):
            raise ValueError(f"Blocked call in Python attachment: {node.func.id}")
        if isinstance(node, ast.Attribute) and node.attr.startswith("__"):
            raise ValueError("Dunder attribute access is blocked in Python attachments")


def execute_python_file(path: str | Path, timeout: int = 15) -> str:
    """Run a validated Python file in isolated interpreter mode and return stdout."""
    file_path = Path(path).resolve()
    if file_path.suffix.lower() != ".py" or not file_path.is_file():
        raise ValueError("Expected an existing .py file")
    source = file_path.read_text(encoding="utf-8")
    _validate(source)
    completed = subprocess.run(
        [
            sys.executable,
            "-I",
            "-c",
            (
                "import runpy,sys,time; "
                "time.sleep=lambda _seconds: None; "
                "runpy.run_path(sys.argv[1],run_name='__main__')"
            ),
            str(file_path),
        ],
        cwd=file_path.parent,
        capture_output=True,
        text=True,
        timeout=timeout,
        check=False,
        env={"PYTHONIOENCODING": "utf-8"},
    )
    if completed.returncode != 0:
        error = completed.stderr.strip()[-2000:]
        raise RuntimeError(f"Python attachment exited {completed.returncode}: {error}")
    output = completed.stdout.strip()
    if not output:
        raise RuntimeError("Python attachment produced no stdout")
    return output