import ast
import builtins
import code
import codeop
import collections.abc
import contextlib
import inspect
import io
import json
import linecache
import logging
import os
import shutil
import subprocess
import sys
import tempfile
import typing
from collections.abc import MutableMapping
from types import CodeType
from typing import Any
from RestrictedPython import (
Eval,
Guards,
RestrictingNodeTransformer,
compile_restricted,
safe_globals,
)
from RestrictedPython.PrintCollector import PrintCollector
from effectful.handlers.llm.template import Tool
from effectful.ops.syntax import ObjectInterpretation, defop, implements
from effectful.ops.types import Operation
[docs]
@defop
def parse(source: str, filename: str) -> ast.Module:
"""
Parse source text into an AST.
source: The Python source code to parse.
filename: The filename recorded in the resulting AST for tracebacks and tooling.
Returns the parsed AST.
"""
raise NotImplementedError(
"An eval provider must be installed in order to parse code."
)
[docs]
@defop
def type_check(
source: str,
lo: int | None = None,
hi: int | None = None,
*,
lenient: bool = False,
) -> None:
"""
Type check a module source, reporting only diagnostics inside a line region.
source: A complete module source to check (e.g. produced by
``splice_into_source``, which splices generated code into a Template's real
module source).
lo, hi: Inclusive line range within ``source`` to report errors from; when
omitted, the whole source is in scope. Errors outside the region are
ignored so unrelated pre-existing code never blocks synthesis.
lenient: when True, relax mypy for incrementally-built REPL code spliced into a
Template body -- allow redefinition (a cell may rebind or redefine a name)
and don't require the body to satisfy the Template's return type. Off (strict)
for synthesized ``Callable`` bodies, which must honor their signature.
Returns None, raises TypeError on an in-region failure.
"""
raise NotImplementedError(
"An eval provider must be installed in order to type check code."
)
[docs]
@defop
def compile(module: ast.Module, filename: str) -> CodeType:
"""
Compile an AST into a Python code object.
module: The AST to compile (typically produced by parse()).
filename: The filename recorded in the resulting code object (CodeType.co_filename), used in tracebacks and by inspect.getsource().
Returns the compiled code object.
"""
raise NotImplementedError(
"An eval provider must be installed in order to compile code."
)
[docs]
@defop
def exec(
bytecode: CodeType,
env: dict[str, Any],
) -> None:
"""
Execute a compiled code object.
bytecode: A code object to execute (typically produced by compile()).
env: The namespace mapping used during execution.
After ``exec(bytecode, env)`` returns, ``env`` reflects all top-level
binding effects of the executed code (new names and rebindings alike).
"""
raise NotImplementedError(
"An eval provider must be installed in order to execute code."
)
logger = logging.getLogger(__name__)
[docs]
def scan_non_nestable(generated: ast.Module) -> None:
"""Reject constructs legal at module level but illegal once nested in a function.
``from ... import *`` and ``from __future__ import ...`` are both ``SyntaxError``s
inside a function body, but mypy *accepts* a nested star import silently, so the
splice would slip an illegal construct past the type check and fail later at
``compile``/``exec``. Detect them explicitly and raise before splicing. Raises
``ValueError`` (this is rejecting invalid generated *source*, not signaling a type
error), so a decoder can catch it alongside ``SyntaxError`` without swallowing a real
``TypeError`` from a broken provider.
"""
for stmt in generated.body:
if isinstance(stmt, ast.ImportFrom):
if stmt.module == "__future__":
raise ValueError(
"generated code uses `from __future__ import ...`, which is "
"illegal once spliced into a function body"
)
if any(alias.name == "*" for alias in stmt.names):
raise ValueError(
"generated code uses a star import (`from ... import *`), which "
"is illegal once spliced into a function body"
)
def _def_nodes(
module: ast.Module,
) -> list[ast.FunctionDef | ast.AsyncFunctionDef]:
"""All function definitions in ``module``, in a stable order that an
``ast.unparse`` -> ``ast.parse`` round-trip preserves (so a def keeps its
index across it)."""
return [
n
for n in ast.walk(module)
if isinstance(n, ast.FunctionDef | ast.AsyncFunctionDef)
]
def _find_def_at_lineno(
module: ast.Module, lineno: int
) -> ast.FunctionDef | ast.AsyncFunctionDef | None:
"""Locate the function definition whose definition site is ``lineno``.
Matches ``fn.__code__.co_firstlineno`` -- the first decorator line, or the
``def`` line when undecorated -- which identifies the def directly and
unambiguously (no name matching, and nesting-agnostic). Returns None only if
no def starts there: a dynamically generated ``fn`` with no source def, or
source that has drifted since import.
"""
for node in _def_nodes(module):
start = node.decorator_list[0].lineno if node.decorator_list else node.lineno
if start == lineno:
return node
return None
def _region_errors(stdout: str, lo: int | None, hi: int | None) -> list[dict[str, Any]]:
"""mypy ``--output=json`` diagnostics of severity ``error`` whose reported
line falls within ``[lo, hi]`` -- the spliced region. An open bound (``None``)
is unbounded on that side, so ``lo=hi=None`` reports every error.
``--output=json`` emits one JSON object per diagnostic carrying mypy's own
``severity`` and ``line`` fields, so we filter on those directly rather than
parsing (and risking mis-parsing) its human-readable format. Only reached
for exit status < 2; a fatal status emits text, not JSON, and is handled by
the caller before this runs.
"""
errors: list[dict[str, Any]] = []
for line in stdout.splitlines():
if not line.strip():
continue
diag = json.loads(line)
if diag["severity"] != "error":
continue
if (lo is None or lo <= diag["line"]) and (hi is None or diag["line"] <= hi):
errors.append(diag)
return errors
[docs]
def splice_into_source(
generated: ast.Module, anchor: Any
) -> tuple[str, int, int] | None:
"""Splice `generated` into the anchor Template's own function body, in its real
module source.
Returns the modified module source and the ``[lo, hi]`` line span of the
spliced body within it, or ``None`` when the anchor's source can't be recovered
(the caller skips rather than guesses). Raises ``RuntimeError`` if the source is
recovered but the anchor's def can't be located in it (source drift) -- a real
error, not a silent pass.
The generated function -- and any helpers it defines alongside -- becomes the
body of the Template's own function at its real (possibly nested) position, so
the generated code is checked in its real lexical scope with no synthesized
type stubs.
"""
if not generated.body:
raise TypeError("splice: generated module is empty")
last = generated.body[-1]
if not isinstance(last, ast.FunctionDef | ast.AsyncFunctionDef):
raise TypeError(
f"splice: last statement must be a function definition, "
f"got {type(last).__name__}"
)
target_name = last.name
recovered = _recover_template_def(anchor)
if recovered is None:
return None
module_ast, template_def = recovered
# Splice in place: replace the body with the generated body and bind the
# target against the (source) return annotation via `return`. Decorators are
# left untouched -- mypy checks a function's body against its declared return
# type regardless of decorators (even an unresolvable / `Any` one), and the
# decorator application itself doesn't spuriously fail, so touching the
# surrounding source as little as possible keeps the splice robust.
template_def.body = [
*generated.body,
ast.Return(ast.Name(target_name, ast.Load())),
]
# mypy reports line numbers in the coordinates of `checked_source`, so we need
# the spliced *body's* span there. ast.unparse reassigns line numbers but
# preserves def order, so the def keeps its index in walk order -- take the def
# at that same index in the re-parsed source.
#
# The region is the body (the generated code) only, NOT the def header: the
# signature and decorators are the Template author's own pre-existing source,
# which we must not attribute to synthesis. This matters for templates whose
# module source can't be fully recovered -- notably notebook/REPL cells, which
# share a runtime namespace but whose recovered source is a single cell missing
# the other cells' imports, so the signature's own annotations (e.g. `Literal`,
# `Callable`) look undefined to mypy. Flagging only the body keeps those
# spurious signature-line diagnostics out of the gate.
def_index = _def_nodes(module_ast).index(template_def)
checked_source = ast.unparse(ast.fix_missing_locations(module_ast))
spliced = _def_nodes(ast.parse(checked_source))[def_index]
lo = spliced.body[0].lineno # first generated statement (body is non-empty)
hi = spliced.end_lineno or lo
return checked_source, lo, hi
def _recover_template_def(
anchor: Any,
) -> tuple[ast.Module, ast.FunctionDef | ast.AsyncFunctionDef] | None:
"""Locate the anchor Template's own ``def`` in its real module source.
Returns the parsed module AST and the def node, or ``None`` when the source can't
be recovered (REPL/exec/notebook Template with no linecache entry -- the caller
skips rather than guesses). Raises ``RuntimeError`` on source drift (source
recovered but the def no longer sits where ``fn`` was compiled from).
"""
fn = inspect.unwrap(anchor) # staticmethod/classmethod -> underlying function
# Recover the module source via fn's own filename -- a real path or a
# linecache-registered synthetic name (e.g. <synthesis:...>) for REPL/exec/
# notebook templates; linecache.getlines reads real files from disk too.
try:
source_file = inspect.getsourcefile(fn)
except TypeError:
source_file = None
module_source = "".join(linecache.getlines(source_file)) if source_file else ""
if not module_source:
logger.warning("skipping type check: cannot recover source for %r", fn)
return None
module_ast = ast.parse(module_source)
template_def = _find_def_at_lineno(module_ast, fn.__code__.co_firstlineno)
if template_def is None:
raise RuntimeError(
f"cannot locate {getattr(fn, '__qualname__', fn)!r} in its module "
f"source (source drifted since import?)"
)
return module_ast, template_def
def _splice_repl(
prior: list[str], snippet: str, anchor: Any
) -> tuple[str, int, int] | None:
"""Splice the cumulative REPL code -- ``prior`` snippets followed by the current
``snippet`` -- into the anchor Template's body, in its real module source, and return
the modified source with the ``[lo, hi]`` line span of the *current* snippet.
The REPL code becomes the Template function's body at its real (possibly nested)
position, so the Template's parameters and enclosing scope -- i.e. the session's seed
env -- are in scope, and each snippet sees the ones before it (they are function
locals). No ``return`` is appended; the REPL code doesn't produce the Template's
declared type, and that contract is waived by ``lenient`` type checking. Every prior
snippet stays in the body so its bindings resolve (matching the runtime, which ran
them), but only the current snippet's lines are reported, so an earlier cell's error
isn't re-reported on every later call.
Returns ``None`` when the current snippet has no statements to check, or when the
Template's source can't be recovered -- a Template defined at a REPL, in a notebook, or
via ``exec()`` is sourceless, so we skip the check and run the code unchecked, exactly
as ``splice_into_source`` does for a sourceless Callable anchor. Raises ``RuntimeError``
only on source *drift* (source recovered but the def no longer sits where it was
compiled from), which ``_recover_template_def`` surfaces.
"""
# An empty or comment-only snippet parses to zero statements: nothing to check.
n_current = len(ast.parse(snippet).body)
if n_current == 0:
return None
# None means the Template's source can't be recovered (REPL/exec/notebook-defined) --
# skip, like the Callable path, rather than break the tool; `_recover_template_def`
# raises on source drift, which is a real error and propagates.
recovered = _recover_template_def(anchor)
if recovered is None:
return None
module_ast, template_def = recovered
cumulative = "".join(s if s.endswith("\n") else s + "\n" for s in [*prior, snippet])
template_def.body = ast.parse(cumulative).body
# mypy reports line numbers in the coordinates of the unparsed source; the current
# snippet is the last `n_current` statements of the spliced body. ast.unparse keeps def
# order, so the template def is at the same walk index after the round-trip.
def_index = _def_nodes(module_ast).index(template_def)
checked_source = ast.unparse(ast.fix_missing_locations(module_ast))
spliced = _def_nodes(ast.parse(checked_source))[def_index]
lo = spliced.body[-n_current].lineno
hi = spliced.body[-1].end_lineno or lo
return checked_source, lo, hi
def _mypy_check_region(
source: str,
lo: int | None = None,
hi: int | None = None,
lenient: bool = False,
) -> None:
"""Run mypy on `source` and raise ``TypeError`` if any error diagnostic falls
within ``[lo, hi]``; raise ``RuntimeError`` if mypy itself fails to run.
Applies mypy to whatever source it's given -- spliced or otherwise -- and
reports only the region's errors (the whole source when the region is
omitted), so pre-existing errors elsewhere in `source` never block synthesis.
When ``lenient`` (for REPL code spliced into a Template body): allow a variable to be
redefined with a new type across cells (``--allow-redefinition``), a def/class/import
to be redefined (``no-redef``), and the body not to return the Template's declared type
(``return``/``empty-body``). All normal for an incrementally-built REPL, not real errors.
"""
lenient_flags = (
[
"--allow-redefinition",
"--disable-error-code=no-redef",
"--disable-error-code=return",
"--disable-error-code=empty-body",
]
if lenient
else []
)
# Run mypy as a subprocess, not the in-process `mypy.api.run`: the API builds
# typeshed and a full module graph inside this process and never returns that
# memory, so under a test/agent session doing many checks it accumulates to many
# GB (OOM). A subprocess reclaims all of it on exit. Pass a file (not --command:
# it hits an argv length limit on large modules); each call gets an isolated temp
# dir + cache so parallel decodes don't share -- and deadlock on -- mypy's cache.
tmpdir = tempfile.mkdtemp(prefix="effectful_typecheck_")
try:
tf_path = os.path.join(tmpdir, "_synthesized.py")
with open(tf_path, "w", encoding="utf-8") as f:
f.write(source)
proc = subprocess.run(
[
sys.executable,
"-m",
"mypy",
tf_path,
"--cache-dir",
os.path.join(tmpdir, "cache"),
"--no-error-summary",
"--output=json",
"--ignore-missing-imports",
"--disable-error-code=import-untyped",
*lenient_flags,
],
capture_output=True,
text=True,
)
stdout, stderr, status = proc.stdout, proc.stderr, proc.returncode
finally:
shutil.rmtree(tmpdir, ignore_errors=True)
# Exit status >= 2 means mypy itself failed (fatal/usage/internal/syntax) -- a
# tool failure, not a type error -- and it emits text rather than JSON, so
# raise `RuntimeError` rather than parse or silently pass.
if status >= 2:
raise RuntimeError(
f"mypy could not check the source:\n{(stdout or '') + (stderr or '')}"
)
errors = _region_errors(stdout or "", lo, hi)
if errors:
# Not the source: it's large and the model already has the generated code.
report = "\n".join(json.dumps(e) for e in errors)
raise TypeError("mypy type check failed:\n" + report)
# Eval Providers
[docs]
class UnsafeEvalProvider(ObjectInterpretation):
"""UNSAFE provider that handles parse, comple and exec operations
by shelling out to python *without* any further checks. Only use for testing."""
[docs]
@implements(type_check)
def type_check(
self,
source: str,
lo: int | None = None,
hi: int | None = None,
*,
lenient: bool = False,
) -> None:
_mypy_check_region(source, lo, hi, lenient)
[docs]
@implements(parse)
def parse(self, source: str, filename: str) -> ast.Module:
# Cache source under `filename` so inspect.getsource() can retrieve it later.
# inspect uses f.__code__.co_filename -> linecache.getlines(filename)
linecache.cache[filename] = (
len(source),
None,
source.splitlines(True),
filename,
)
return ast.parse(source, filename=filename, mode="exec")
[docs]
@implements(compile)
def compile(self, module: ast.AST, filename: str) -> CodeType:
return builtins.compile(typing.cast(typing.Any, module), filename, "exec")
[docs]
@implements(exec)
def exec(
self,
bytecode: CodeType,
env: dict[str, Any],
) -> None:
# Ensure builtins exist in the execution environment.
env.setdefault("__builtins__", __builtins__)
# Execute module-style so top-level defs land in `env`.
builtins.exec(bytecode, env, env)
class _StdoutPrintCollector(PrintCollector):
"""`_print_` factory whose `print(...)` writes to the real `sys.stdout`
(so output-capturing callers see it) rather than accumulating into the
collector's discarded `printed` buffer."""
def _call_print(self, *objects, **kwargs):
kwargs.setdefault("file", sys.stdout)
builtins.print(*objects, **kwargs)
[docs]
class RestrictedEvalProvider(ObjectInterpretation):
"""
Safer provider using RestrictedPython.
RestrictedPython is not a complete sandbox, but it enforces a restricted
language subset and expects you to provide a constrained exec environment.
policy : dict[str, Any], optional
RestrictedPython compile_restricted policy for compilation
"""
policy: type[RestrictingNodeTransformer] | None = None
def __init__(
self,
*,
policy: type[RestrictingNodeTransformer] | None = None,
):
self.policy = policy
[docs]
@implements(type_check)
def type_check(
self,
source: str,
lo: int | None = None,
hi: int | None = None,
*,
lenient: bool = False,
) -> None:
_mypy_check_region(source, lo, hi, lenient)
[docs]
@implements(parse)
def parse(self, source: str, filename: str) -> ast.Module:
# Keep inspect.getsource() working for dynamically-defined objects.
linecache.cache[filename] = (
len(source),
None,
source.splitlines(True),
filename,
)
return ast.parse(source, filename=filename, mode="exec")
[docs]
@implements(compile)
def compile(self, module: ast.Module, filename: str) -> CodeType:
# RestrictedPython can compile from an AST directly.
return compile_restricted(
module,
filename=filename,
mode="exec",
policy=self.policy or RestrictingNodeTransformer,
)
[docs]
@implements(exec)
def exec(
self,
bytecode: CodeType,
env: dict[str, Any],
) -> None:
# Build restricted globals from RestrictedPython's defaults
rglobals: dict[str, Any] = safe_globals.copy()
# Enable class definitions (required for Python 3)
rglobals["__metaclass__"] = type
rglobals["__name__"] = "restricted"
# Layer `env` on top (without letting callers replace the restricted builtins).
rglobals.update({k: v for k, v in env.items() if k != "__builtins__"})
# Enable for loops and comprehensions
rglobals["_getiter_"] = Eval.default_guarded_getiter
# Enable sequence unpacking in comprehensions and for loops
rglobals["_iter_unpack_sequence_"] = Guards.guarded_iter_unpack_sequence
rglobals["getattr"] = Guards.safer_getattr
rglobals["setattr"] = Guards.guarded_setattr
rglobals["_write_"] = lambda x: x
# RestrictedPython rewrites `print(...)` into its `_print_` collector
# protocol; route it to the real stdout so output-capturing callers
# (e.g. redirect_stdout) see it instead of a discarded collector.
rglobals["_print_"] = _StdoutPrintCollector
# Snapshot value identities before execution so we can copy back every
# *binding effect* — both new names and rebindings of seeded names.
before = dict(rglobals)
builtins.exec(bytecode, rglobals, rglobals)
sentinel = object()
env.update(
{
key: value
for key, value in rglobals.items()
if key != "__builtins__" and before.get(key, sentinel) is not value
}
)
class _OpCommandCompiler(codeop.CommandCompiler):
"""A `codeop.CommandCompiler` that routes compilation through the
`parse`/`compile` effect operations (so the installed eval provider owns it
and `parse` populates `linecache`), replacing the native single-mode
compiler that `code.InteractiveInterpreter` installs.
"""
def __call__(
self, source: str, filename: str = "<input>", symbol: str = "single"
) -> CodeType:
# `runsource` passes symbol="single"; we ignore it and compile in the
# exec mode the ops produce, so a complete multi-statement block runs in
# one shot. Incomplete/invalid input raises SyntaxError, which
# `runsource` routes to `showsyntaxerror` (we do not buffer partial input
# -- there is no line-at-a-time protocol).
return compile(parse(source, filename), filename)
[docs]
class ReplSession(code.InteractiveInterpreter):
"""A persistent, output-capturing Python session seeded from a lexical
context.
`exec_code(source)` runs a pre-compiled code object in `self.locals` through
the `exec` effect operation. Both bindings and captured stdout/stderr
persist across calls -- variables, imports and definitions accumulate exactly
like a REPL -- and the session (with its buffer) is discarded as a whole when
it goes out of scope. Each call returns only the output it produced; a
snippet that raises has its traceback appended to that output rather than
propagating -- mirroring `code.InteractiveInterpreter`, only `SystemExit`
propagates -- so failures are surfaced as text. There is no bare-expression
auto-echo, so use `print()` to surface values.
Compilation -- and therefore syntax checking -- happens earlier, at the
`Encodable[CodeType]` boundary; this session only executes.
"""
# The session's captured output, accumulated across calls and exposed for
# introspection. stdout (`print` output) and stderr (writes plus tracebacks)
# are kept separate; `exec_code` returns each call's slice of both.
stdout: io.StringIO
stderr: io.StringIO
def __init__(self, env: MutableMapping[str, Any]):
# Run in a fresh writable dict seeded with a flat view of `env`. This is
# forced by `exec`: its globals must be one real dict (a ChainMap is
# rejected), and a REPL needs a single persistent namespace so a function
# defined in one snippet sees a name a later snippet binds. Seeding a flat
# copy also leaves the lexical seed untouched, so REPL assignments never
# leak into the surrounding scope.
scope: dict[str, Any] = dict(env)
# When `env` is the per-call `ChainMap` (its outer layers are read-only
# frame proxies), splice this dict in as an extra shadowing first layer so
# the bindings are *also* visible to the rest of the Template call
# (mirroring `exec`) -- still scoped to the call, since that ChainMap is.
if isinstance(env, collections.ChainMap):
env.maps.insert(0, scope)
# `InteractiveInterpreter.__init__` stores it as `self.locals`, so we reuse
# the base's runcode/showtraceback/write machinery.
super().__init__(scope)
# Route `runsource`'s compilation through the `parse`/`compile` ops too, so
# it stays consistent with our `runcode` (which execs through the `exec`
# op) rather than the native single-mode compiler the base installed.
self.compile = _OpCommandCompiler()
self.stdout = io.StringIO()
self.stderr = io.StringIO()
self._prior_snippets: list[str] = []
@property
def prior_snippets(self) -> list[str]:
"""Sources of the actual error-free executed snippets, in order -- the type-check
context the `Encodable[CodeType]` decoder splices before the current snippet."""
return self._prior_snippets
[docs]
def runcode(self, code: CodeType) -> None:
# Mirrors `InteractiveInterpreter.runcode` exactly; the only difference
# is that `exec` here is the effect operation, so execution routes
# through the installed eval provider. `showtraceback` reports failures
# via `self.write`, which `exec_code` has redirected into `self.stderr`.
try:
exec(code, self.locals)
except SystemExit:
raise
except:
self.showtraceback()
[docs]
@Tool.define
def exec_code(self, code: CodeType) -> str:
"""Run Python in a persistent, stateful session and return its output.
This is a long-lived REPL, not a one-shot sandbox: every call runs in the
SAME namespace, so names you bind in one call stay available in later
calls within the same task. Imports, function/class definitions and
variable assignments all accumulate during the session of this template.
The namespace starts seeded with the in-scope variables of the surrounding context, which you may read and
rebind.
Output: returns this call's output -- its stdout (what `print` wrote)
followed by its stderr (which includes the traceback if the code raised).
There is NO automatic echoing of results -- a bare expression on its own
line (e.g. `1 + 1`) displays nothing, so call `print(...)` for anything
you want to see. A snippet that raises has its traceback returned and the
session survives, so you can read the error and continue in the next call
(only `SystemExit` aborts).
Provide `code` as a string of Python source. It must be a complete,
compilable snippet -- incomplete or invalid source is rejected before it
runs.
"""
out_start = self.stdout.tell()
err_start = self.stderr.tell()
# Record this snippet's source so the *next* snippet's decode-time type check can
# splice the accumulated session code into the Template body. The type check itself
# lives in the `Encodable[CodeType]` decoder (as it does for synthesized Callables),
# not here -- this session only runs code.
self._prior_snippets.append("".join(linecache.getlines(code.co_filename)))
with (
contextlib.redirect_stdout(self.stdout),
contextlib.redirect_stderr(self.stderr),
):
self.runcode(code)
return self.stdout.getvalue()[out_start:] + self.stderr.getvalue()[err_start:]
@Operation.define
def _repl_session(env: MutableMapping[str, Any]) -> "ReplSession":
"""Return the REPL session for the current Template call, seeded from `env`.
`PythonRepl` (in completions.py) installs a fresh handler for this inside each
`Template.__apply__` (mirroring how `__history__` is managed), giving the session a
lifetime of exactly one Template call. Outside such a scope there is no managed
session, so this falls back to a fresh one -- e.g. when tools are listed outside a
Template call, or when a code object is decoded with no REPL in scope.
Defined here (not with `PythonRepl`) so the `Encodable[CodeType]` decoder can reach the
session -- and its accumulated `prior_snippets` -- at decode time without importing
`completions` (which would be a cycle).
"""
return ReplSession(env)