|
| 1 | +"""Adapt the DSPy facade to Monty's native Python subset. |
| 2 | +
|
| 3 | +Classes, methods, closures, and calls belong to Monty. We only remove the shim's |
| 4 | +no-op Module base/initializer and supply its attribute and subscription hooks. |
| 5 | +Validate first so unsupported object-model features fail before guest execution. |
| 6 | +""" |
| 7 | + |
| 8 | +import ast |
| 9 | +from dataclasses import dataclass |
| 10 | + |
| 11 | +from dspy.primitives.code_interpreter import CodeExecutionError |
| 12 | + |
| 13 | + |
| 14 | +def _name(name, ctx=None): |
| 15 | + return ast.Name(id=name, ctx=ctx or ast.Load()) |
| 16 | + |
| 17 | + |
| 18 | +def _call(name, *args): |
| 19 | + return ast.Call(func=_name(name), args=list(args), keywords=[]) |
| 20 | + |
| 21 | + |
| 22 | +def _unsupported(node, message): |
| 23 | + raise CodeExecutionError(f"Unsupported Monty syntax at sandbox:{node.lineno}: {message}") |
| 24 | + |
| 25 | + |
| 26 | +def _module_init(node): |
| 27 | + return ( |
| 28 | + isinstance(node, ast.Call) |
| 29 | + and isinstance(node.func, ast.Attribute) |
| 30 | + and node.func.attr == "__init__" |
| 31 | + and isinstance(node.func.value, ast.Call) |
| 32 | + and isinstance(node.func.value.func, ast.Name) |
| 33 | + and node.func.value.func.id == "super" |
| 34 | + and not node.func.value.args |
| 35 | + and not node.func.value.keywords |
| 36 | + ) |
| 37 | + |
| 38 | + |
| 39 | +class _Validate(ast.NodeVisitor): |
| 40 | + """Reject semantics we cannot preserve before executing any generated code.""" |
| 41 | + |
| 42 | + _dynamic = frozenset({"eval", "exec", "compile", "globals", "locals", "vars", "dir", "__import__"}) |
| 43 | + _unsupported_nodes = frozenset( |
| 44 | + { |
| 45 | + "Delete", |
| 46 | + "Yield", |
| 47 | + "YieldFrom", |
| 48 | + "Match", |
| 49 | + "TryStar", |
| 50 | + "AsyncWith", |
| 51 | + "AsyncFor", |
| 52 | + "AsyncFunctionDef", |
| 53 | + } |
| 54 | + ) |
| 55 | + |
| 56 | + def __init__(self): |
| 57 | + self.in_function = False |
| 58 | + self.module_inits = set() |
| 59 | + self.direct_method = None |
| 60 | + |
| 61 | + def generic_visit(self, node): |
| 62 | + if type(node).__name__ in self._unsupported_nodes: |
| 63 | + _unsupported(node, type(node).__name__) |
| 64 | + super().generic_visit(node) |
| 65 | + |
| 66 | + def _identifier(self, node, name): |
| 67 | + if name.startswith(("_dspy", "_Dspy", "__dspy")): |
| 68 | + _unsupported(node, "the _dspy / _Dspy / __dspy namespaces are reserved for the guest runtime") |
| 69 | + if name in {"super", "__class__"}: |
| 70 | + _unsupported(node, "only direct super().__init__(...) in the module initializer is supported") |
| 71 | + |
| 72 | + def visit_Name(self, node): |
| 73 | + self._identifier(node, node.id) |
| 74 | + if node.id in self._dynamic: |
| 75 | + _unsupported(node, f"{node.id} exposes or replaces compiled execution scope") |
| 76 | + |
| 77 | + def visit_arg(self, node): |
| 78 | + self._identifier(node, node.arg) |
| 79 | + self.generic_visit(node) |
| 80 | + |
| 81 | + def visit_alias(self, node): |
| 82 | + self._identifier(node, node.asname or node.name.split(".")[0]) |
| 83 | + if node.name == "*": |
| 84 | + _unsupported(node, "wildcard imports") |
| 85 | + |
| 86 | + def visit_ClassDef(self, node): |
| 87 | + self._identifier(node, node.name) |
| 88 | + if self.in_function: |
| 89 | + _unsupported(node, "nested class definitions") |
| 90 | + if node.decorator_list or node.keywords or len(node.bases) != 1: |
| 91 | + _unsupported(node, "use one dspy.Module base, without class decorators or metaclasses") |
| 92 | + for item in node.body: |
| 93 | + if isinstance(item, ast.FunctionDef): |
| 94 | + if item.decorator_list: |
| 95 | + _unsupported(item, "method decorators and descriptors") |
| 96 | + if item.name.startswith("__") and item.name != "__init__": |
| 97 | + _unsupported(item, "custom special methods; call forward explicitly") |
| 98 | + if item.name == "__init__": |
| 99 | + self.module_inits.update( |
| 100 | + stmt.value for stmt in item.body if isinstance(stmt, ast.Expr) and _module_init(stmt.value) |
| 101 | + ) |
| 102 | + elif not ( |
| 103 | + isinstance(item, ast.Pass) or (isinstance(item, ast.Expr) and isinstance(item.value, ast.Constant)) |
| 104 | + ): |
| 105 | + _unsupported(item, "class bodies may contain only methods and docstrings") |
| 106 | + self.generic_visit(node) |
| 107 | + |
| 108 | + def visit_FunctionDef(self, node): |
| 109 | + self._identifier(node, node.name) |
| 110 | + previous = self.in_function |
| 111 | + self.in_function = True |
| 112 | + self.generic_visit(node) |
| 113 | + self.in_function = previous |
| 114 | + |
| 115 | + def visit_Attribute(self, node): |
| 116 | + if node.attr.startswith("__") and not node.attr.endswith("__"): |
| 117 | + _unsupported(node, "name-mangled private attributes; use a single underscore") |
| 118 | + # Only diagnose receivers whose type follows from syntax. Inferring |
| 119 | + # names from assignments or annotations would misclassify rebinding, |
| 120 | + # branch joins, supplied tools, and compiled Flex methods. |
| 121 | + literal_type = { |
| 122 | + ast.List: list, |
| 123 | + ast.ListComp: list, |
| 124 | + ast.Tuple: tuple, |
| 125 | + ast.Dict: dict, |
| 126 | + ast.DictComp: dict, |
| 127 | + ast.Set: set, |
| 128 | + ast.SetComp: set, |
| 129 | + ast.JoinedStr: str, |
| 130 | + }.get(type(node.value)) |
| 131 | + if isinstance(node.value, ast.Constant): |
| 132 | + literal_type = type(node.value.value) |
| 133 | + if ( |
| 134 | + node is not self.direct_method |
| 135 | + and isinstance(node.ctx, ast.Load) |
| 136 | + and literal_type is not None |
| 137 | + and callable(getattr(literal_type, node.attr, None)) |
| 138 | + ): |
| 139 | + _unsupported( |
| 140 | + node, |
| 141 | + f"native method '{node.attr}' used as a value. Call it directly, or use a named helper: " |
| 142 | + f"def helper(value, *args, **kwargs): return value.{node.attr}(*args, **kwargs)", |
| 143 | + ) |
| 144 | + self.generic_visit(node) |
| 145 | + |
| 146 | + def visit_Call(self, node): |
| 147 | + if node in self.module_inits: |
| 148 | + for argument in [*node.args, *node.keywords]: |
| 149 | + self.visit(argument) |
| 150 | + return |
| 151 | + previous = self.direct_method |
| 152 | + self.direct_method = node.func |
| 153 | + self.generic_visit(node) |
| 154 | + self.direct_method = previous |
| 155 | + |
| 156 | + def visit_comprehension(self, node): |
| 157 | + if any(isinstance(n, (ast.Attribute, ast.Subscript)) for n in ast.walk(node.target)): |
| 158 | + _unsupported(node.target, "attribute/subscript targets in comprehensions; use an explicit loop") |
| 159 | + if node.is_async: |
| 160 | + _unsupported(node.target, "async comprehensions") |
| 161 | + self.generic_visit(node) |
| 162 | + |
| 163 | + |
| 164 | +class _Compile(ast.NodeTransformer): |
| 165 | + def __init__(self): |
| 166 | + self.augments = 0 |
| 167 | + |
| 168 | + def visit(self, node): |
| 169 | + result = super().visit(node) |
| 170 | + if hasattr(node, "lineno"): |
| 171 | + for item in result if isinstance(result, list) else [result]: |
| 172 | + if isinstance(item, ast.AST): |
| 173 | + ast.copy_location(item, node) |
| 174 | + return result |
| 175 | + |
| 176 | + def visit_Import(self, node): |
| 177 | + return [ |
| 178 | + ast.Assign(targets=[_name(alias.asname or "dspy", ast.Store())], value=_name("_dspy_namespace")) |
| 179 | + if alias.name == "dspy" |
| 180 | + else ast.Import(names=[alias]) |
| 181 | + for alias in node.names |
| 182 | + ] |
| 183 | + |
| 184 | + def visit_ImportFrom(self, node): |
| 185 | + if node.module != "dspy" or node.level: |
| 186 | + return node |
| 187 | + return [ |
| 188 | + ast.Assign( |
| 189 | + targets=[_name(alias.asname or alias.name, ast.Store())], |
| 190 | + value=_call("_dspy_getattr", _name("_dspy_namespace"), ast.Constant(alias.name)), |
| 191 | + ) |
| 192 | + for alias in node.names |
| 193 | + ] |
| 194 | + |
| 195 | + def visit_ClassDef(self, node): |
| 196 | + base = self.visit(node.bases[0]) |
| 197 | + node.bases = [] |
| 198 | + node = self.generic_visit(node) |
| 199 | + return [ |
| 200 | + ast.Expr(value=_call("_dspy_check_base", base)), |
| 201 | + node, |
| 202 | + ast.Expr( |
| 203 | + value=ast.Call( |
| 204 | + func=ast.Attribute(value=_name("_dspy_module_types"), attr="append", ctx=ast.Load()), |
| 205 | + args=[_name(node.name)], |
| 206 | + keywords=[], |
| 207 | + ) |
| 208 | + ), |
| 209 | + ] |
| 210 | + |
| 211 | + def visit_Attribute(self, node): |
| 212 | + node.value = self.visit(node.value) |
| 213 | + if isinstance(node.ctx, ast.Load): |
| 214 | + return _call("_dspy_getattr", node.value, ast.Constant(node.attr)) |
| 215 | + return node |
| 216 | + |
| 217 | + def visit_Call(self, node): |
| 218 | + if _module_init(node): |
| 219 | + node.func = _name("_dspy_module_init") |
| 220 | + elif isinstance(node.func, ast.Attribute): |
| 221 | + # Native method calls work in Monty even when fetching that method |
| 222 | + # with getattr does not. Do not route calls through an emulator. |
| 223 | + node.func.value = self.visit(node.func.value) |
| 224 | + else: |
| 225 | + node.func = self.visit(node.func) |
| 226 | + node.args = [self.visit(arg) for arg in node.args] |
| 227 | + node.keywords = [self.visit(kw) for kw in node.keywords] |
| 228 | + return node |
| 229 | + |
| 230 | + def key(self, value): |
| 231 | + if isinstance(value, ast.Slice): |
| 232 | + return _call( |
| 233 | + "_dspy_slice", |
| 234 | + *(self.visit(part) if part else ast.Constant(None) for part in (value.lower, value.upper, value.step)), |
| 235 | + ) |
| 236 | + if isinstance(value, ast.Tuple): |
| 237 | + return ast.Tuple(elts=[self.key(item) for item in value.elts], ctx=ast.Load()) |
| 238 | + return self.visit(value) |
| 239 | + |
| 240 | + def visit_Subscript(self, node): |
| 241 | + if isinstance(node.ctx, ast.Load): |
| 242 | + return _call("_dspy_getitem", self.visit(node.value), self.key(node.slice)) |
| 243 | + return self.generic_visit(node) |
| 244 | + |
| 245 | + def visit_AugAssign(self, node): |
| 246 | + if not isinstance(node.target, ast.Attribute): |
| 247 | + return self.generic_visit(node) |
| 248 | + # A prediction field needs the same read hook in `p.x += rhs` as in |
| 249 | + # `p.x`. Evaluate the receiver once, and preserve in-place mutation. |
| 250 | + self.augments += 1 |
| 251 | + obj, value = f"_dspy_obj_{self.augments}", f"_dspy_value_{self.augments}" |
| 252 | + return [ |
| 253 | + ast.Assign(targets=[_name(obj, ast.Store())], value=self.visit(node.target.value)), |
| 254 | + ast.Assign( |
| 255 | + targets=[_name(value, ast.Store())], |
| 256 | + value=_call("_dspy_getattr", _name(obj), ast.Constant(node.target.attr)), |
| 257 | + ), |
| 258 | + ast.AugAssign(target=_name(value, ast.Store()), op=node.op, value=self.visit(node.value)), |
| 259 | + ast.Assign( |
| 260 | + targets=[ast.Attribute(value=_name(obj), attr=node.target.attr, ctx=ast.Store())], |
| 261 | + value=_name(value), |
| 262 | + ), |
| 263 | + ] |
| 264 | + |
| 265 | + |
| 266 | +@dataclass |
| 267 | +class CompiledCode: |
| 268 | + source: str |
| 269 | + original: str |
| 270 | + lines: dict[int, int] |
| 271 | + |
| 272 | + def annotate(self, error): |
| 273 | + """Attach original statement locations, without depending on feed numbers.""" |
| 274 | + generated, original = self.source.splitlines(), self.original.splitlines() |
| 275 | + locations = [] |
| 276 | + for frame in error.traceback(): |
| 277 | + line = frame.line |
| 278 | + if line in self.lines and frame.source_line.strip() == generated[line - 1].strip(): |
| 279 | + source_line = self.lines[line] |
| 280 | + locations.append(f" sandbox:{source_line}: {original[source_line - 1].strip()}") |
| 281 | + return CodeExecutionError("\n".join([str(error), *locations])) |
| 282 | + |
| 283 | + |
| 284 | +def compile_source(source: str) -> CompiledCode: |
| 285 | + tree = ast.parse(source) |
| 286 | + _Validate().visit(tree) |
| 287 | + tree = ast.fix_missing_locations(_Compile().visit(tree)) |
| 288 | + rendered = ast.unparse(tree) |
| 289 | + # Pair statements, not expression nodes: unparse may normalize expressions |
| 290 | + # but preserves statement structure. Inner statement spans override outer |
| 291 | + # spans, yielding a source map without relying on ast's private unparser. |
| 292 | + before = [node for node in ast.walk(tree) if isinstance(node, ast.stmt)] |
| 293 | + after = [node for node in ast.walk(ast.parse(rendered)) if isinstance(node, ast.stmt)] |
| 294 | + lines = {} |
| 295 | + for old, new in zip(before, after, strict=True): |
| 296 | + assert type(old) is type(new) |
| 297 | + for line in range(new.lineno, new.end_lineno + 1): |
| 298 | + lines[line] = old.lineno |
| 299 | + return CompiledCode(rendered, source, lines) |
0 commit comments