From 914170dea3dff263aa9018ac07a9ca12cdf0b330 Mon Sep 17 00:00:00 2001 From: Eli Date: Sat, 25 Jul 2026 18:53:32 -0400 Subject: [PATCH 1/3] Don't route all-scalar monoid ops through jax `_jax_args` admitted `jax.typing.ArrayLike`, a union that includes `bool`, `int`, `float` and `complex`, so the jax `Monoid.plus` handlers claimed pure-Python scalar arithmetic. They extend `EvaluateIntp` after the scalar implementations and so take precedence, silently narrowing a Python float to a `float32` array and leaving downstream rules treating a scalar body as array-valued. Require at least one genuine array. Co-Authored-By: Claude Opus 5 (1M context) --- effectful/handlers/jax/monoid.py | 19 ++++++++---- tests/test_handlers_jax_monoid.py | 49 +++++++++++++++++++++++++++++++ 2 files changed, 63 insertions(+), 5 deletions(-) diff --git a/effectful/handlers/jax/monoid.py b/effectful/handlers/jax/monoid.py index 3acf89f88..2e99d6859 100644 --- a/effectful/handlers/jax/monoid.py +++ b/effectful/handlers/jax/monoid.py @@ -68,12 +68,21 @@ def _jax_args(args): - """True iff ``args`` is non-empty and every arg is a concrete - :class:`jax.typing.ArrayLike` or named tensor. - + """True iff ``args`` is non-empty, every arg is a concrete + :class:`jax.typing.ArrayLike` or named tensor, and at least one of them is + an array rather than a Python scalar. + + :class:`jax.typing.ArrayLike` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so admitting it alone would claim pure-Python + scalar arithmetic. These handlers extend ``EvaluateIntp`` after the scalar + implementations and so take precedence over them, which would silently + narrow a Python float to a ``float32`` array and leave downstream rules + treating a scalar body as array-valued. """ - return args and all( - isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args + return ( + args + and all(isinstance(a, jax.typing.ArrayLike) or is_eager_array(a) for a in args) + and any(not isinstance(a, bool | int | float | complex) for a in args) ) diff --git a/tests/test_handlers_jax_monoid.py b/tests/test_handlers_jax_monoid.py index 14a1db08e..a538664d8 100644 --- a/tests/test_handlers_jax_monoid.py +++ b/tests/test_handlers_jax_monoid.py @@ -124,6 +124,55 @@ def test_reduce_array_2(monoid, reductor, backend: JaxBackend): assert jnp.allclose(actual, expected) +SCALAR_PLUS = [ + pytest.param(Sum, 3.0, id="Sum"), + pytest.param(Product, 2.0, id="Product"), +] + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_scalars_stays_scalar(monoid, expected): + """``Monoid.plus`` of plain Python numbers must not become a ``jax.Array``. + + ``jax.typing.ArrayLike`` is a union that includes ``bool``, ``int``, + ``float`` and ``complex``, so the jax ``plus`` handlers -- which extend + ``EvaluateIntp`` after the scalar implementations and therefore take + precedence -- must not claim pure-Python scalar arithmetic. + """ + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(1.0, 2.0) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize( + "monoid,expected", + [pytest.param(Sum, 6.0, id="Sum"), pytest.param(Product, 8.0, id="Product")], +) +def test_reduce_scalar_body_stays_scalar(monoid, expected, backend: JaxBackend): + """A reduction over a scalar body is likewise plain Python arithmetic.""" + i = backend.define_vars("i", ret="scalar") + + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.reduce(2.0, {i: range(3)}) + + assert not isinstance(actual, jax.Array) + assert isinstance(actual, float) + assert actual == expected + + +@pytest.mark.parametrize("monoid,expected", SCALAR_PLUS) +def test_plus_mixed_array_and_scalar_is_array(monoid, expected): + """One genuine array is enough: the narrowing must not be over-applied.""" + with handler(NormalizeIntp), handler(EvaluateIntp): + actual = monoid.plus(jnp.asarray([1.0, 1.0]), 2.0) + + assert isinstance(actual, jax.Array) + assert jnp.allclose(actual, jnp.asarray([expected, expected])) + + @pytest.mark.parametrize("monoid,reductor", MONOIDS) def test_arange_reduce_indirect(monoid, reductor, backend: JaxBackend): """When the range var is used both as a direct index and as a value From f9cbeaff9a1c076984a1506d5cdff8c9903dfda0 Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 28 Jul 2026 11:07:52 -0400 Subject: [PATCH 2/3] Add a generator-expression bytecode disassembler `effectful/internals/disassembly.py` symbolically interprets the bytecode of a generator expression (and of lambdas and comprehensions nested inside it) back into an `ast` node, so a comprehension's source syntax can be recovered from the code object at runtime. Supports CPython 3.12 and 3.13. Standalone: imports nothing from `effectful` and touches no existing code. Co-Authored-By: Claude Opus 5 (1M context) --- effectful/internals/disassembly.py | 3438 ++++++++++++++++++++++++++ tests/test_internals_disassembler.py | 1764 +++++++++++++ 2 files changed, 5202 insertions(+) create mode 100644 effectful/internals/disassembly.py create mode 100644 tests/test_internals_disassembler.py diff --git a/effectful/internals/disassembly.py b/effectful/internals/disassembly.py new file mode 100644 index 000000000..47aaccb52 --- /dev/null +++ b/effectful/internals/disassembly.py @@ -0,0 +1,3438 @@ +""" +Generator expression bytecode reconstruction module. + +This module provides functionality to reconstruct AST representations from compiled +generator expressions by analyzing their bytecode. The primary use case is to recover +the original structure of generator comprehensions from their compiled form. + +The only public-facing interface is the `disassemble()` function, which takes a +generator object and returns an AST node representing the original comprehension. +All other functions and classes in this module are internal implementation details. + +Example: + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> # ast_node is now an ast.Expression representing the original expression +""" + +import ast +import builtins +import collections +import collections.abc +import copy +import dis +import enum +import functools +import inspect +import itertools +import sys +import types +import typing +from collections.abc import Callable, Generator, Iterator +from dataclasses import dataclass, field, replace + +CompExp = ast.GeneratorExp | ast.ListComp | ast.SetComp | ast.DictComp + + +class Placeholder(ast.Name): + """Placeholder for AST nodes that are not yet resolved.""" + + def __init__( + self, + id: typing.Literal[".PLACEHOLDER"] = ".PLACEHOLDER", + ctx: ast.Load = ast.Load(), + ): + super().__init__(id=id, ctx=ctx) + + +class DummyIterName(ast.Name): + """Dummy name for the iterator variable in generator expressions.""" + + def __init__(self, id: typing.Literal[".0"] = ".0", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class Skipped(ast.Name): + """Placeholder for skipped branches in if-expressions. + + ``id`` is defaulted so that ``copy.deepcopy`` can reconstruct the node: on + Python 3.12 ``ast.AST.__reduce__`` supplies no positional arguments (3.13+ + supplies one per field), so the constructor must be callable with none. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class CommonConstant(ast.Name): + """A constant pushed by 3.14's LOAD_COMMON_CONSTANT. + + 3.14 can inline `any()`/`all()` over a generator, guarding the fast path + with `loaded_name is `. Marking the builtin lets that guard be + recognised so the generic call path is followed instead, which still spells + out the original call. + """ + + def __init__(self, id: str = "", ctx: ast.Load = ast.Load()): + super().__init__(id=id, ctx=ctx) + + +class TargetHole(ast.Name): + """Placeholder for a comprehension loop target that is not yet named. + + ``FOR_ITER`` knows that a loop target exists but not what it is called; the + name only arrives with the ``STORE_FAST``/``UNPACK_SEQUENCE`` instructions + that follow. Each hole carries a unique ``id`` so that it can be located + again inside a comprehension after the surrounding state has been copied, + which matters for unpacking targets where several holes are live at once. + """ + + _counter: typing.ClassVar[Iterator[int]] = itertools.count() + + def __init__(self, id: str = "", ctx: ast.Store = ast.Store()): + super().__init__(id=id or f".TARGET_{next(TargetHole._counter)}", ctx=ctx) + + +class ReplaceTargetHole(ast.NodeTransformer): + """Replace the uniquely-identified :class:`TargetHole` ``id`` with ``replacement``.""" + + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = replacement + super().__init__() + + def visit_TargetHole(self, node: TargetHole) -> ast.expr: + return self.replacement if node.id == self.id else node + + +def _bind_target_hole( + stack: list[ast.expr], hole: ast.expr, replacement: ast.expr +) -> list[ast.expr]: + """Fill the loop-target hole ``hole`` of the innermost matching comprehension. + + Returns a new stack; the hole itself is left in place for the caller to pop. + """ + assert isinstance(hole, TargetHole), f"Expected a loop target hole, got {hole}" + for pos, item in zip(reversed(range(len(stack))), reversed(stack)): + if not isinstance(item, CompExp) or not item.generators: + continue + if not any( + isinstance(n, TargetHole) and n.id == hole.id + for n in ast.walk(item.generators[-1].target) + ): + continue + new_comp = ReplaceTargetHole(hole.id, replacement).visit(copy.deepcopy(item)) + return stack[:pos] + [new_comp] + stack[pos + 1 :] + + raise TypeError(f"No comprehension found with loop target hole {hole.id}") + + +class Null(ast.Constant): + """Placeholder for NULL values generated in bytecode.""" + + def __init__(self, value: None = None): + super().__init__(value=value) + + +class ConvertedValue(ast.expr): + """Wrapper for values that have been converted with CONVERT_VALUE.""" + + value: ast.expr + conversion: int + ast_conversion: int + + def __init__(self, value: ast.expr, conversion: int): + self.value = value + self.conversion = conversion + # Map CONVERT_VALUE args to ast.FormattedValue conversion values + # CONVERT_VALUE: 0=None, 1=str, 2=repr, 3=ascii + # ast.FormattedValue: -1=none, 115=str, 114=repr, 97=ascii + conversion_map = {0: -1, 1: 115, 2: 114, 3: 97} + self.ast_conversion = conversion_map.get(conversion, -1) + + +class CompLambda(ast.Lambda): + """Placeholder AST node representing a lambda function used in comprehensions.""" + + def __init__(self, body: CompExp): + assert isinstance(body, CompExp) + assert sum(1 for x in ast.walk(body) if isinstance(x, DummyIterName)) == 1 + assert len(body.generators) > 0 + assert isinstance(body.generators[0].iter, DummyIterName) + args = ast.arguments( + posonlyargs=[ast.arg(DummyIterName().id)], + args=[], + kwonlyargs=[], + kw_defaults=[], + defaults=[], + ) + super().__init__(args=args, body=body) + + def __copy__(self): + """Support copy.copy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(self.body) + + def __deepcopy__(self, memo): + """Support copy.deepcopy operation.""" + assert isinstance(self.body, CompExp) + return CompLambda(copy.deepcopy(self.body, memo)) + + def inline(self, iterator: ast.expr) -> CompExp: + assert isinstance(self.body, CompExp) + res: CompExp = copy.deepcopy(self.body) + res.generators[0].iter = iterator + return res + + +class ReplacePlaceholder(ast.NodeTransformer): + value: ast.expr + _done: bool + + def __init__(self, value: ast.expr): + self.value = value + self._done = False + super().__init__() + + def visit(self, node): + if isinstance(node, Placeholder) and not self._done: + self._done = True + return self.value + else: + return self.generic_visit(node) + + +class ReplaceSkipped(ast.NodeTransformer): + id: str + replacement: ast.expr + + def __init__(self, id: str, replacement: ast.expr): + self.id = id + self.replacement = copy.deepcopy(replacement) + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped) and node.body.id == self.id: + return ast.IfExp(test=node.test, body=self.replacement, orelse=node.orelse) + elif isinstance(node.orelse, Skipped) and node.orelse.id == self.id: + return ast.IfExp(test=node.test, body=node.body, orelse=self.replacement) + else: + return self.generic_visit(node) + + +class BranchState(typing.NamedTuple): + testval: bool + value: ast.expr + + +class BranchIdentifier(ast.NodeVisitor): + branching: collections.abc.MutableMapping[str, BranchState] + filter_lengths: list[int] + + def __init__(self): + self.branching = {} + self.filter_lengths = [] + super().__init__() + + def visit_IfExp(self, node: ast.IfExp): + if isinstance(node.body, Skipped): + self.branching[node.body.id] = BranchState( + testval=False, value=copy.deepcopy(node.orelse) + ) + elif isinstance(node.orelse, Skipped): + self.branching[node.orelse.id] = BranchState( + testval=True, value=copy.deepcopy(node.body) + ) + return self.generic_visit(node) + + def visit_comprehension(self, node: ast.comprehension): + self.filter_lengths.append(len(node.ifs)) + return self.generic_visit(node) + + +@functools.cache +def _instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Decode a code object once; every state derived from it shares the result.""" + return collections.OrderedDict( + (instr.offset, instr) for instr in dis.get_instructions(code) + ) + + +@functools.cache +def _next_instructions( + code: types.CodeType, +) -> collections.abc.Mapping[int, dis.Instruction]: + """Map each instruction offset to the instruction that follows it.""" + ordered = list(_instructions(code).values()) + return {before.offset: after for before, after in zip(ordered[:-1], ordered[1:])} + + +@dataclass(frozen=True) +class ReconstructionState: + """State maintained during AST reconstruction from bytecode. + + This class tracks all the information needed while processing bytecode + instructions to reconstruct the original comprehension's AST. It acts + as the working memory during the reconstruction process, maintaining + both the evaluation stack state and the high-level comprehension structure + being built. + + The reconstruction process works by simulating the Python VM's execution + of the bytecode, but instead of executing operations, it builds AST nodes + that represent those operations. + + Attributes: + code: The compiled code object from which the bytecode is being processed. + This is typically obtained from a generator function or comprehension. + + stack: Simulates the Python VM's value stack. Contains AST nodes or + values that would be on the stack during execution. Operations + like LOAD_FAST push to this stack, while operations like + BINARY_ADD pop operands and push results. + """ + + code: types.CodeType + instruction: dis.Instruction + + stack: list[ast.expr] = field(default_factory=list) + result: ast.expr = field(default_factory=Placeholder) + + # How many times each FOR_ITER has been entered on this path. + loops: dict[int, int] = field(default_factory=dict) + finished: bool = field(default=False) + + # Which edge each already-resolved conditional jump took on this path. + branches: "dict[int, BranchEdge]" = field(default_factory=dict) + + # Locals bound to a known expression rather than to a loop target. 3.14 + # unrolls a single-iteration loop over a literal, storing its targets + # directly, so those names have to be substituted back at their uses. + bindings: dict[str, ast.expr] = field(default_factory=dict) + + # Set by KW_NAMES (Python 3.12 only) and consumed by the following CALL. + # KW_NAMES has no stack effect, so the names cannot live on `stack`. + kw_names: tuple[str, ...] | None = field(default=None) + + @property + def instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + """The bytecode instructions of the current code object, by offset.""" + return _instructions(self.code) + + @property + def next_instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: + return _next_instructions(self.code) + + +# Python version enum for version-specific handling +class PythonVersion(enum.IntEnum): + PY_312 = 12 + PY_313 = 13 + PY_314 = 14 + + +def current_version() -> PythonVersion: + """The bytecode dialect of the running interpreter. + + Raises on a Python this module has not been taught, rather than guessing + that the previous release's opcodes still mean what they used to. + """ + try: + return PythonVersion(sys.version_info.minor) + except ValueError as e: + supported = ", ".join(f"3.{v.value}" for v in PythonVersion) + raise NotImplementedError( + f"effectful.internals.disassembly supports {supported}, " + f"not 3.{sys.version_info.minor}" + ) from e + + +# Global handler registry +OpHandler = Callable[[ReconstructionState, dis.Instruction], ReconstructionState] + +OP_HANDLERS: dict[str, OpHandler] = {} + + +@typing.overload +def register_handler( + opname: str, *, version: PythonVersion +) -> Callable[[OpHandler], OpHandler]: ... + + +@typing.overload +def register_handler( + opname: str, + handler: OpHandler, + *, + version: PythonVersion, +) -> OpHandler: ... + + +def register_handler( + opname: str, + handler=None, + *, + version: PythonVersion, +): + """Register a handler for one opcode in one Python bytecode dialect. + + Every dialect a handler applies to is named explicitly. Opcodes are not + assumed to carry forward: a release can keep an opcode's name while changing + what it does, so applicability to a new Python is a decision to make per + opcode rather than a default. + """ + if handler is None: + return functools.partial(register_handler, opname, version=version) + + # Skip registration if version doesn't match current version + if version != current_version(): + return handler + + # Only check opmap if the version matches (or no version specified) + assert opname in dis.opmap, f"Invalid operation name: '{opname}'" + + if opname in OP_HANDLERS: + raise ValueError(f"Handler for '{opname}' (version {version}) already exists.") + + if dis.opmap[opname] in dis.hasjrel: + assert opname in LOOP_OPS | BRANCH_OPS | JUMP_OPS + else: + assert opname not in LOOP_OPS | BRANCH_OPS | JUMP_OPS + + @functools.wraps(handler) + def _wrapper( + state: ReconstructionState, + instr: dis.Instruction, + ) -> ReconstructionState: + assert instr.opname == opname, ( + f"Handler for '{opname}' called with wrong instruction" + ) + assert not state.finished, "Cannot process instruction on finished state" + + new_state = handler(state, instr) + + jump: bool | None # argument to dis.stack_effect + if instr.opname in LOOP_OPS: + if state.loops.get(instr.offset, 0) > 0: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + else: + # Copy rather than mutate: continuations forked from this state + # share the mapping and must not see each other's loop counts. + new_state = replace( + new_state, + instruction=state.next_instructions[instr.offset], + loops={ + **state.loops, + instr.offset: state.loops.get(instr.offset, 0) + 1, + }, + ) + jump = False + elif instr.opname in BRANCH_OPS: + if new_state.branches.get(instr.offset) == BranchEdge.FALL_THROUGH: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = False + else: + new_state = replace( + new_state, instruction=state.instructions[instr.argval] + ) + jump = True + elif instr.opname in JUMP_OPS: + new_state = replace(new_state, instruction=state.instructions[instr.argval]) + jump = True + elif instr.opname not in RETURN_OPS and instr.offset in state.next_instructions: + new_state = replace( + new_state, instruction=state.next_instructions[instr.offset] + ) + jump = None + else: + new_state = replace(new_state, finished=True) + jump = None + + # post-condition: check stack effect + expected_stack_effect = dis.stack_effect(instr.opcode, instr.arg, jump=jump) + actual_stack_effect = len(new_state.stack) - len(state.stack) + assert len(state.stack) + expected_stack_effect >= 0, ( + f"Handler for '{opname}' would result in negative stack size" + ) + assert actual_stack_effect == expected_stack_effect, ( + f"Handler for '{opname}' has incorrect stack effect: " + f"expected {expected_stack_effect}, got {actual_stack_effect}" + ) + + return new_state + + OP_HANDLERS[opname] = _wrapper + return handler # return the original handler for multiple decorator usage + + +LOOP_OPS: set[typing.Literal["FOR_ITER"]] = {"FOR_ITER"} + +BRANCH_OPS: set[ + typing.Literal[ + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", + ] +] = { + "POP_JUMP_IF_TRUE", + "POP_JUMP_IF_FALSE", + "POP_JUMP_IF_NOT_NONE", + "POP_JUMP_IF_NONE", +} + +RETURN_OPS: set[typing.Literal["RETURN_VALUE", "RETURN_CONST"]] = { + "RETURN_VALUE", + "RETURN_CONST", +} + +JUMP_OPS = {dis.opname[d] for d in dis.hasjrel} - LOOP_OPS - BRANCH_OPS - RETURN_OPS + + +# Instructions that emit an element of the comprehension being built. Reaching +# one of these means the current iteration was *not* filtered out. +PRODUCE_OPS = {"YIELD_VALUE", "LIST_APPEND", "SET_ADD", "MAP_ADD"} + + +def _successor_offsets(state: ReconstructionState, instr: dis.Instruction) -> list[int]: + """Offsets control can transfer to from ``instr``, ignoring exception edges.""" + following = state.next_instructions.get(instr.offset) + if instr.opname in BRANCH_OPS | LOOP_OPS: + return [instr.argval] + ([following.offset] if following else []) + elif instr.opname in JUMP_OPS: + return [instr.argval] + elif instr.opname in RETURN_OPS: + return [] + else: + return [following.offset] if following else [] + + +def _reachable_outcomes(state: ReconstructionState, start: int) -> tuple[bool, bool]: + """From ``start``, can the iteration be skipped, and can an element be produced? + + Returns ``(can_skip, can_produce)``. "Skip" means reaching the loop + back-edge without emitting an element, i.e. being filtered out. + """ + seen: set[int] = set() + pending = [start] + can_skip = can_produce = False + + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + can_produce = True + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + can_skip = True + continue + + pending.extend(_successor_offsets(state, instr)) + + return can_skip, can_produce + + +class BranchEdge(enum.IntEnum): + """Which way a conditional jump was resolved on the path being explored.""" + + TAKE_JUMP = 1 + FALL_THROUGH = 2 + + +class BranchKind(enum.Enum): + """What role a conditional jump plays in a comprehension. + + TERNARY + A conditional expression: the arms reconverge having each pushed a + value, and both are spliced back together into an ``ast.IfExp``. + FILTER + Part of a filter's condition. The condition is consumed rather than + producing a value, so each surviving path records the conjunction of + tests that got it to the element, and the filter as a whole is the + disjunction of those conjunctions. + """ + + TERNARY = enum.auto() + FILTER = enum.auto() + + +@functools.cache +def _stack_depths(code: types.CodeType) -> collections.abc.Mapping[int, int]: + """VM stack depth on entry to each reachable instruction. + + Depths are relative to the start of the code object, which is all that is + needed to tell a value-producing branch from a control-flow one. + """ + instructions = collections.OrderedDict( + (i.offset, i) for i in dis.get_instructions(code) + ) + ordered = list(instructions.values()) + following = {a.offset: b.offset for a, b in zip(ordered[:-1], ordered[1:])} + + depths: dict[int, int] = {ordered[0].offset: 0} + pending = collections.deque([ordered[0].offset]) + while pending: + offset = pending.popleft() + instr, depth = instructions[offset], depths[offset] + if instr.opname in RETURN_OPS: + continue + + edges: list[tuple[int, bool | None]] = [] + if instr.opname in BRANCH_OPS | LOOP_OPS: + edges = [(instr.argval, True)] + if offset in following: + edges.append((following[offset], False)) + elif instr.opname in JUMP_OPS: + edges = [(instr.argval, True)] + elif offset in following: + edges = [(following[offset], None)] + + for target, jump in edges: + if target in instructions and target not in depths: + depths[target] = depth + dis.stack_effect( + instr.opcode, instr.arg, jump=jump + ) + pending.append(target) + + return depths + + +def _forward_reachable(state: ReconstructionState, start: int) -> set[int]: + """Offsets reachable from ``start`` without producing or looping back.""" + seen: set[int] = set() + pending = [start] + while pending: + offset = pending.pop() + if offset in seen or offset not in state.instructions: + continue + seen.add(offset) + + instr = state.instructions[offset] + if instr.opname in PRODUCE_OPS: + continue + if ( + instr.opname == "JUMP_BACKWARD" + and state.instructions[instr.argval].opname in LOOP_OPS + ): + continue + + pending.extend(_successor_offsets(state, instr)) + + return seen + + +def _is_conditional_expression( + state: ReconstructionState, instr: dis.Instruction +) -> bool: + """Do the two edges of ``instr`` reconverge one stack slot deeper? + + That is the signature of a conditional expression: each arm leaves a value + behind and control rejoins to consume it. A filter's condition is consumed + by the jump itself, so wherever its edges meet -- if they meet at all -- the + stack is no deeper than it was. + """ + following = state.next_instructions.get(instr.offset) + if following is None: + return False + + common = _forward_reachable(state, instr.argval) & _forward_reachable( + state, following.offset + ) + if not common: + return False # the edges never rejoin, so nothing was left on the stack + + depths = _stack_depths(state.code) + join = min(common) # arms are laid out contiguously, so the join comes first + if join not in depths or following.offset not in depths: + return False + return depths[join] == depths[following.offset] + 1 + + +def _classify_branch( + state: ReconstructionState, instr: dis.Instruction +) -> tuple[BranchKind, list[BranchEdge]]: + """Classify a conditional jump and list the edges worth exploring.""" + both = [BranchEdge.TAKE_JUMP, BranchEdge.FALL_THROUGH] + following = state.next_instructions.get(instr.offset) + if following is None: + return BranchKind.TERNARY, both + + jump_skip, _ = _reachable_outcomes(state, instr.argval) + fall_skip, _ = _reachable_outcomes(state, following.offset) + + # Neither edge can drop the current iteration -- because there is no loop at + # all (a lambda body) or because the element is produced regardless (a + # conditional in the element expression). Either way nothing is filtered. + if not jump_skip and not fall_skip: + return BranchKind.TERNARY, both + + # Otherwise the branch could be a filter, or a conditional expression that + # merely happens to sit inside one. Only the latter leaves a value behind. + if _is_conditional_expression(state, instr): + return BranchKind.TERNARY, both + + # An edge that cannot reach an element contributes nothing to the filter, so + # there is no point walking it. Pruning those edges is also what keeps the + # executor out of the operand-cleanup block on a chained comparison's + # failing edge. + live = [ + edge + for edge, start in ( + (BranchEdge.TAKE_JUMP, instr.argval), + (BranchEdge.FALL_THROUGH, following.offset), + ) + if _reachable_outcomes(state, start)[1] + ] + return BranchKind.FILTER, live or [BranchEdge.TAKE_JUMP] + + +def _negate(condition: ast.expr) -> ast.expr: + """Logical negation, cancelling a `not` rather than stacking another one.""" + if isinstance(condition, ast.UnaryOp) and isinstance(condition.op, ast.Not): + return condition.operand + return ast.UnaryOp(op=ast.Not(), operand=condition) + + +def _conjoin(conditions: list[ast.expr]) -> ast.expr | None: + """Combine the entries of a ``comprehension.ifs`` list into one expression.""" + if not conditions: + return None + elif len(conditions) == 1: + return conditions[0] + else: + return ast.BoolOp(op=ast.And(), values=list(conditions)) + + +def _disjoin(left: ast.expr, right: ast.expr) -> ast.expr: + """Combine two conditions with ``or``, flattening nested disjunctions. + + Two rewrites are applied while combining. Duplicate disjuncts are dropped, + because paths through independent filters repeat them. And ``X or (not X and + Y)`` becomes ``X or Y``: enumerating paths records the negation of every + test a path declined, so a later disjunct restates the negation of an + earlier one. Dropping it is not merely tidier -- `or` short-circuits, so the + earlier disjunct has already been evaluated, and leaving the negation in + would evaluate it a second time, which is visibly wrong when the condition + contains an assignment expression. + + Conditions are keyed by ``ast.dump`` exactly once each: these lists get long + and the expressions large, so re-dumping per comparison dominates. + """ + values: list[ast.expr] = [] + for side in (left, right): + if isinstance(side, ast.BoolOp) and isinstance(side.op, ast.Or): + values.extend(side.values) + else: + values.append(side) + + unique: list[ast.expr] = [] + seen: set[str] = set() + for value in values: + key = ast.dump(value) + if key in seen: + continue + seen.add(key) + + # Absorb the negations of the disjuncts already accepted. + if isinstance(value, ast.BoolOp) and isinstance(value.op, ast.And): + kept = [ + conjunct + for conjunct in value.values + if not ( + isinstance(conjunct, ast.UnaryOp) + and isinstance(conjunct.op, ast.Not) + and ast.dump(conjunct.operand) in seen + ) + ] + if kept and len(kept) < len(value.values): + conjoined = _conjoin(kept) + assert conjoined is not None + value = conjoined + + unique.append(value) + + return unique[0] if len(unique) == 1 else ast.BoolOp(op=ast.Or(), values=unique) + + +def _merge_filters_into( + node: typing.Any, other: typing.Any, mutate: bool = True +) -> bool: + """Walk two results in parallel, OR-ing the filters where they disagree. + + Everything outside a ``comprehension.ifs`` has to match exactly; the ifs are + where the paths are allowed to differ, and are combined rather than + compared. Filters are not recursed into, so a nested comprehension inside a + filter is treated as part of that filter's condition. + + Returns False if the two results differ somewhere they may not. With + ``mutate=False`` nothing is written, which allows compatibility to be tested + before paying for a deep copy -- most candidate pairs do not merge, and the + copy dominates otherwise. + """ + if type(node) is not type(other): + return False + + if isinstance(node, ast.comprehension): + if ast.dump(node.target) != ast.dump(other.target): + return False + if not _merge_filters_into(node.iter, other.iter, mutate): + return False + if not mutate: + return True + + guard, other_guard = _conjoin(node.ifs), _conjoin(other.ifs) + if guard is None or other_guard is None: + # One path reached the element unconditionally, so the filter as a + # whole is unconditional at this generator. + node.ifs = [] + elif ast.dump(guard) != ast.dump(other_guard): + node.ifs = [_disjoin(guard, other_guard)] + return True + + if not isinstance(node, ast.AST): + return bool(node == other) + + for name in node._fields: + mine, theirs = getattr(node, name, None), getattr(other, name, None) + if isinstance(mine, list) or isinstance(theirs, list): + if not isinstance(mine, list) or not isinstance(theirs, list): + return False + if len(mine) != len(theirs): + return False + if not all(_merge_filters_into(a, b, mutate) for a, b in zip(mine, theirs)): + return False + elif isinstance(mine, ast.AST) or isinstance(theirs, ast.AST): + if not isinstance(mine, ast.AST) or not isinstance(theirs, ast.AST): + return False + if not _merge_filters_into(mine, theirs, mutate): + return False + elif mine != theirs: + return False + + return True + + +def _merge_filters(left: ast.expr, right: ast.expr) -> ast.expr | None: + """Combine two paths that differ only in which filter conditions they met. + + Each path through a filter records the conjunction that got it to the + element; the filter as a whole is the disjunction over all such paths. + Returns ``None`` when the results differ by more than their filters. + """ + # A marker anywhere means some conditional expression is still unresolved, + # and unresolved arms must be spliced before anything can be OR-ed. + if any(isinstance(n, Skipped) for n in ast.walk(left)): + return None + if any(isinstance(n, Skipped) for n in ast.walk(right)): + return None + + if not _merge_filters_into(left, right, mutate=False): + return None + + merged = copy.deepcopy(left) + return merged if _merge_filters_into(merged, right) else None + + +def _skipped_offset(key: str) -> int: + """Sort key for `.SKIPPED_` markers, so merging is deterministic.""" + return int(key.rsplit("_", 1)[-1]) + + +def _merge_at_ifexp(left: ast.expr, right: ast.expr) -> ast.expr: + """ + Merge two expression ASTs obtained from two branches of symbolic execution. + """ + if isinstance(left, ast.Constant) and left.value is None: + return copy.deepcopy(right) + elif isinstance(right, ast.Constant) and right.value is None: + return copy.deepcopy(left) + + assert type(left) == type(right) + + lb, rb = BranchIdentifier(), BranchIdentifier() + lb.visit(left) + rb.visit(right) + + # A conditional expression: each path filled in one arm and left a marker in + # the other, so splice the two together. Sorted for determinism -- set + # iteration order over the marker names varies with PYTHONHASHSEED. + common_keys = set(lb.branching) & set(rb.branching) + differing = [ + key + for key in sorted(common_keys, key=_skipped_offset) + if lb.branching[key].testval != rb.branching[key].testval + ] + + # Only copy once it is known there is something to splice; this runs for + # every candidate pair of paths, most of which have nothing in common. + merged: ast.expr = copy.deepcopy(left) if differing else left + for key in differing: + visited = ReplaceSkipped(key, rb.branching[key].value).visit(merged) + assert isinstance(visited, ast.expr) + merged = visited + spliced = bool(differing) + + # The paths may *also* have satisfied different filter conditions on the way + # to the element, so combine those too rather than picking one arbitrarily. + combined = _merge_filters(merged, right) + if combined is not None: + return combined + if spliced: + return merged + + if ast.dump(left) == ast.dump(right): + return copy.deepcopy(left) + + raise ValueError("No differing branches found to merge") + + +def _specialization_guard_edge( + state: ReconstructionState, instr: dis.Instruction +) -> BranchEdge | None: + """The edge past a 3.14 inlined-builtin guard, or None if this isn't one. + + 3.14 may inline `any()`/`all()` over a generator, guarding the inlined code + with `loaded_name is ` and keeping an ordinary call on the + other edge. Only that other edge still contains the call to reconstruct, so + the guard is treated as though the identity test failed. + """ + if not state.stack: + return None + condition = state.stack[-1] + if not isinstance(condition, ast.Compare): + return None + if not any(isinstance(c, CommonConstant) for c in condition.comparators): + return None + + # Follow the edge taken when the identity test is false. + if instr.opname == "POP_JUMP_IF_FALSE": + return BranchEdge.TAKE_JUMP + elif instr.opname == "POP_JUMP_IF_TRUE": + return BranchEdge.FALL_THROUGH + else: + return None + + +def _merge_all(results: list[ast.expr]) -> ast.expr: + """Combine every path's result into one expression. + + Merging is not associative: a path that still carries an unfilled + conditional-expression arm can only combine with the path that took the + other arm, which need not be its neighbour. So rather than folding in + order, repeatedly merge whichever pair actually combines. + """ + pending = list(results) + while len(pending) > 1: + for i, j in itertools.combinations(range(len(pending)), 2): + try: + merged = _merge_at_ifexp(pending[i], pending[j]) + except (ValueError, AssertionError): + continue + pending = [merged] + [p for k, p in enumerate(pending) if k not in (i, j)] + break + else: + raise ValueError("Could not merge the paths of symbolic execution") + + return pending[0] + + +def _decide_branch( + state: ReconstructionState, instr: dis.Instruction, edge: BranchEdge +) -> ReconstructionState: + """Record which edge of ``instr`` the path being explored takes.""" + return replace(state, branches={**state.branches, instr.offset: edge}) + + +def _symbolic_exec(code: types.CodeType) -> ast.expr: + """Execute bytecode symbolically, following control flow.""" + continuations: list[ReconstructionState] = [ + ReconstructionState( + code=code, + instruction=next(iter(dis.get_instructions(code))), + stack=[Placeholder(), Placeholder()] + if current_version() == PythonVersion.PY_312 + and code.co_flags & inspect.CO_GENERATOR + else [Placeholder()], + ) + ] + + results: list[ast.expr] = [] + + while continuations: + state = continuations.pop() + while not state.finished: + instr = state.instruction + if instr.opname in BRANCH_OPS and instr.offset not in state.branches: + forced = _specialization_guard_edge(state, instr) + if forced is not None: + state = _decide_branch(state, instr, forced) + else: + _, live = _classify_branch(state, instr) + # Explore the first live edge now; queue the rest for later. + continuations.extend( + _decide_branch(state, instr, edge) for edge in live[1:] + ) + state = _decide_branch(state, instr, live[0]) + + state = OP_HANDLERS[state.instruction.opname](state, state.instruction) + results.append(state.result) + + assert results, "No results from symbolic execution" + result = _merge_all(results) + assert not any(isinstance(n, Skipped) for n in ast.walk(result)), ( + "Every conditional expression arm must have been filled in" + ) + return result + + +# ============================================================================ +# GENERATOR COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_312) +def handle_return_generator_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 2 and all( + isinstance(x, Null | Placeholder) for x in state.stack + ), "RETURN_GENERATOR must be the first instruction" + new_result = ast.GeneratorExp(elt=Placeholder(), generators=[]) + return replace(state, stack=[new_result, Null()]) + + +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_313) +@register_handler("RETURN_GENERATOR", version=PythonVersion.PY_314) +def handle_return_generator( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_GENERATOR is the first instruction in generator expressions in Python 3.13+ + assert len(state.stack) == 1 and isinstance(state.stack[0], Null | Placeholder), ( + "RETURN_GENERATOR must be the first instruction" + ) + return replace( + state, stack=[ast.GeneratorExp(elt=Placeholder(), generators=[]), Null()] + ) + + +@register_handler("YIELD_VALUE", version=PythonVersion.PY_312) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_313) +@register_handler("YIELD_VALUE", version=PythonVersion.PY_314) +def handle_yield_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # YIELD_VALUE pops a value from the stack and yields it + # This is the expression part of the generator + assert isinstance(state.result, Placeholder) + new_result = copy.deepcopy(state.stack[0]) + assert isinstance(new_result, ast.GeneratorExp), ( + "YIELD_VALUE must be called after RETURN_GENERATOR" + ) + assert len(new_result.generators) > 0, "YIELD_VALUE should have generators" + assert any(isinstance(x, Placeholder) for x in ast.walk(new_result.elt)) + new_result.elt = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit( + new_result.elt + ) + new_stack = [new_result] + state.stack[1:] + return replace(state, stack=new_stack, result=new_result) + + +# ============================================================================ +# LIST COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_LIST", version=PythonVersion.PY_312) +@register_handler("BUILD_LIST", version=PythonVersion.PY_313) +@register_handler("BUILD_LIST", version=PythonVersion.PY_314) +def handle_build_list( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + # Check if this looks like the start of a list comprehension pattern + # In nested comprehensions, BUILD_LIST(0) starts a new list comprehe + new_ret = ast.ListComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_ret] + return replace(state, stack=new_stack) + else: + # BUILD_LIST with elements - create a regular list + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.List(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_APPEND", version=PythonVersion.PY_312) +@register_handler("LIST_APPEND", version=PythonVersion.PY_313) +@register_handler("LIST_APPEND", version=PythonVersion.PY_314) +def handle_list_append( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.ListComp) + + # add the body to the comprehension + comp: ast.ListComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# SET COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_SET", version=PythonVersion.PY_312) +@register_handler("BUILD_SET", version=PythonVersion.PY_313) +@register_handler("BUILD_SET", version=PythonVersion.PY_314) +def handle_build_set( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.SetComp(elt=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + elements = [ensure_ast(elem) for elem in state.stack[-size:]] + new_stack = state.stack[:-size] + elt_node = ast.Set(elts=elements) + new_stack = new_stack + [elt_node] + return replace(state, stack=new_stack) + + +@register_handler("SET_ADD", version=PythonVersion.PY_312) +@register_handler("SET_ADD", version=PythonVersion.PY_313) +@register_handler("SET_ADD", version=PythonVersion.PY_314) +def handle_set_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + + # add the body to the comprehension + comp: ast.SetComp = copy.deepcopy(state.stack[-instr.argval - 1]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.elt)) + comp.elt = ReplacePlaceholder(state.stack[-1]).visit(comp.elt) + + # swap the return value + new_stack = state.stack[:-1] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# DICT COMPREHENSION HANDLERS +# ============================================================================ + + +@register_handler("BUILD_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_MAP", version=PythonVersion.PY_313) +@register_handler("BUILD_MAP", version=PythonVersion.PY_314) +def handle_build_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + size: int = instr.arg + + if size == 0: + new_result = ast.DictComp(key=Placeholder(), value=Placeholder(), generators=[]) + new_stack = state.stack + [new_result] + return replace(state, stack=new_stack) + else: + # Pop key-value pairs for the dict + keys: list[ast.expr | None] = [ + ensure_ast(state.stack[-2 * i - 2]) for i in range(size) + ] + values = [ensure_ast(state.stack[-2 * i - 1]) for i in range(size)] + new_stack = state.stack[: -2 * size] if size > 0 else state.stack + + # Create dict AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("MAP_ADD", version=PythonVersion.PY_312) +@register_handler("MAP_ADD", version=PythonVersion.PY_313) +@register_handler("MAP_ADD", version=PythonVersion.PY_314) +def handle_map_add( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.stack[-instr.argval - 2], ast.DictComp) + + # add the body to the comprehension + comp: ast.DictComp = copy.deepcopy(state.stack[-instr.argval - 2]) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.key)) + assert any(isinstance(x, Placeholder) for x in ast.walk(comp.value)) + comp.key = ReplacePlaceholder(state.stack[-2]).visit(comp.key) + comp.value = ReplacePlaceholder(state.stack[-1]).visit(comp.value) + + # swap the return value + new_stack = state.stack[:-2] + new_stack[-instr.argval] = comp + + return replace(state, stack=new_stack) + + +# ============================================================================ +# LOOP CONTROL HANDLERS +# ============================================================================ + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_312) +@register_handler("RETURN_VALUE", version=PythonVersion.PY_313) +def handle_return_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert isinstance(state.result, Placeholder) + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + new_stack = state.stack[:-1] + return replace(state, stack=new_stack, result=new_result) + + +def _unyielded_comprehension(state: ReconstructionState) -> CompExp | None: + """The comprehension of a body the compiler proved unreachable, if any. + + An always-false filter lets the compiler drop the whole body: the loop is + still walked, but nothing is ever yielded or appended, so no element + expression survives. The partly built comprehension still carries its + generators, so it can be rebuilt with a filter that is never satisfied -- + which iterates exactly as the original did and produces nothing. + """ + for item in reversed(state.stack): + if not isinstance(item, CompExp) or not item.generators: + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + if not isinstance(element, Placeholder): + continue + + unreachable = copy.deepcopy(item) + never = ast.Constant(value=None) + if isinstance(unreachable, ast.DictComp): + unreachable.key, unreachable.value = never, copy.deepcopy(never) + else: + unreachable.elt = never + unreachable.generators[-1].ifs = [ast.Constant(value=False)] + return unreachable + + return None + + +@register_handler("RETURN_VALUE", version=PythonVersion.PY_314) +def handle_return_value_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Two things changed in 3.14. RETURN_CONST is gone, so a generator's + # trailing `return None` now arrives as LOAD_CONST + RETURN_VALUE; and + # RETURN_VALUE's stack effect is 0 rather than -1, the returned value being + # discarded along with the frame. The value therefore stays on the stack. + if not isinstance(state.result, Placeholder): + assert ( + isinstance(state.stack[-1], ast.Constant) and state.stack[-1].value is None + ), "A generator may only fall off the end returning None" + return state + + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + + assert len(state.stack) == 2 + new_result = ReplacePlaceholder(ensure_ast(state.stack[-1])).visit(state.stack[-2]) + return replace(state, result=new_result) + + +@register_handler("RETURN_CONST", version=PythonVersion.PY_312) +@register_handler("RETURN_CONST", version=PythonVersion.PY_313) +def handle_return_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RETURN_CONST returns a constant value (replaces some LOAD_CONST + RETURN_VALUE patterns) + # Similar to RETURN_VALUE but with a constant + if isinstance(state.result, Placeholder): + unreachable = _unyielded_comprehension(state) + if unreachable is not None: + return replace(state, result=unreachable) + return replace(state, result=ensure_ast(instr.argval)) + else: + assert instr.argval is None + return state + + +@register_handler("FOR_ITER", version=PythonVersion.PY_312) +@register_handler("FOR_ITER", version=PythonVersion.PY_313) +@register_handler("FOR_ITER", version=PythonVersion.PY_314) +def handle_for_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FOR_ITER pops an iterator from the stack and pushes the next item + # If the iterator is exhausted, it jumps to the target instruction + assert len(state.stack) > 0, "FOR_ITER must have an iterator on the stack" + + if state.loops.get(instr.offset, 0) > 0: + return replace(state, stack=state.stack + [Null()]) + + # The iterator should be on top of stack + iterator: ast.expr = state.stack[-1] + + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + element = item.value if isinstance(item, ast.DictComp) else item.elt + new_result = copy.deepcopy(item) + + if isinstance(element, Placeholder): + loop_iter = ensure_ast(iterator) + elif isinstance(element, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(element) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be this loop's iterable, as in + # `for y in (a if c else b)`. Move it back out and plug the value + # this path produced into the arm still awaiting one. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + plugged = ReplacePlaceholder(ensure_ast(iterator)).visit( + copy.deepcopy(element) + ) + assert isinstance(plugged, ast.expr) + loop_iter = plugged + else: + continue + + # The loop target is not named until the STORE_* that follows. + loop_info = ast.comprehension( + target=TargetHole(), iter=loop_iter, ifs=[], is_async=0 + ) + new_result.generators.append(loop_info) + new_stack = ( + state.stack[:pos] + + [new_result] + + state.stack[pos + 1 :] + + [loop_info.target] + ) + return replace(state, stack=new_stack) + + raise TypeError("FOR_ITER did not find partial comprehension on stack") + + +@register_handler("GET_ITER", version=PythonVersion.PY_312) +@register_handler("GET_ITER", version=PythonVersion.PY_313) +@register_handler("GET_ITER", version=PythonVersion.PY_314) +def handle_get_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # GET_ITER converts the top stack item to an iterator + # For AST reconstruction, we typically don't need to change anything + # since the iterator will be used directly in the comprehension + return state + + +@register_handler("END_FOR", version=PythonVersion.PY_312) +def handle_end_for_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop, followed by POP_TOP (in 3.12) + new_stack = state.stack[:-2] + return replace(state, stack=new_stack) + + +@register_handler("END_FOR", version=PythonVersion.PY_313) +@register_handler("END_FOR", version=PythonVersion.PY_314) +def handle_end_for( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # END_FOR marks the end of a for loop - no action needed for AST reconstruction + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +@register_handler("RERAISE", version=PythonVersion.PY_312) +@register_handler("RERAISE", version=PythonVersion.PY_313) +@register_handler("RERAISE", version=PythonVersion.PY_314) +def handle_reraise( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RERAISE re-raises an exception - generally ignore for AST reconstruction + return state + + +# ============================================================================ +# VARIABLE OPERATIONS HANDLERS +# ============================================================================ + + +def _literal_elements(value: ast.expr, count: int | None = None) -> list[ast.expr]: + """The elements of a literal sequence, for destructuring a known value.""" + assert isinstance(value, ast.Tuple | ast.List), ( + f"Cannot unpack {type(value).__name__}; expected a literal sequence" + ) + assert count is None or len(value.elts) == count, ( + f"Expected {count} values to unpack, got {len(value.elts)}" + ) + return [ensure_ast(element) for element in value.elts] + + +def _bind_local( + state: ReconstructionState, var_name: str, value: ast.expr +) -> ReconstructionState: + """Record that a local now stands for ``value``, popping it off the stack. + + Reached when a store is not filling in a loop target: 3.14 unrolls a + single-iteration loop over a literal, assigning its targets outright. The + loop is gone from the bytecode, so the names it bound are not in scope in + the reconstruction and their uses are substituted instead. + """ + bindings = {**state.bindings, var_name: ensure_ast(value)} + return replace(state, stack=state.stack[:-1], bindings=bindings) + + +def _read_local(state: ReconstructionState, var_name: str) -> ast.expr: + """The expression a local name stands for at this point on this path.""" + if var_name == DummyIterName().id: + return DummyIterName() + elif var_name in state.bindings: + # Bound to a known expression rather than by a loop, so the name itself + # is not in scope in the reconstruction; use what it was bound to. + return copy.deepcopy(state.bindings[var_name]) + else: + return ast.Name(id=var_name, ctx=ast.Load()) + + +@register_handler("LOAD_FAST", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST", version=PythonVersion.PY_314) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_CHECK", version=PythonVersion.PY_314) +def handle_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_CHECK differs only in raising when the local is unbound, which + # says nothing about the expression being reconstructed. + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_DEREF", version=PythonVersion.PY_312) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_313) +@register_handler("LOAD_DEREF", version=PythonVersion.PY_314) +def handle_load_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_DEREF loads a value from a closure variable + var_name = instr.argval + new_stack = state.stack + [ast.Name(id=var_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_312) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_313) +@register_handler("LOAD_CLOSURE", version=PythonVersion.PY_314) +def handle_load_closure( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_CLOSURE loads a closure variable + var_name = instr.argval + new_stack = state.stack + [ast.Name(id=var_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_CONST", version=PythonVersion.PY_312) +@register_handler("LOAD_CONST", version=PythonVersion.PY_313) +@register_handler("LOAD_CONST", version=PythonVersion.PY_314) +def handle_load_const( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + const_value = instr.argval + new_stack = state.stack + [ensure_ast(const_value)] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_312) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_313) +@register_handler("LOAD_GLOBAL", version=PythonVersion.PY_314) +def handle_load_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + global_name = instr.argval + + if instr.argrepr.endswith(" + NULL"): + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load()), Null()] + elif instr.argrepr.startswith("NULL + "): + new_stack = state.stack + [Null(), ast.Name(id=global_name, ctx=ast.Load())] + else: + new_stack = state.stack + [ast.Name(id=global_name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +@register_handler("LOAD_NAME", version=PythonVersion.PY_312) +@register_handler("LOAD_NAME", version=PythonVersion.PY_313) +@register_handler("LOAD_NAME", version=PythonVersion.PY_314) +def handle_load_name( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_NAME is similar to LOAD_GLOBAL but for names in the global namespace + name = instr.argval + new_stack = state.stack + [ast.Name(id=name, ctx=ast.Load())] + return replace(state, stack=new_stack) + + +def _is_assignment_expression(state: ReconstructionState) -> bool: + """Is the top of the stack a COPY of the value beneath it? + + That duplication is how an assignment expression keeps its value after + binding it: `COPY 1` then a `STORE_*`. `handle_copy` pushes the very same + node, so identity is what distinguishes it from two equal-looking values. + """ + return len(state.stack) >= 2 and state.stack[-1] is state.stack[-2] + + +def _handle_assignment_expression( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + """Rebuild `(name := value)` from the COPY/STORE pair that implements it.""" + target = ast.Name(id=instr.argval, ctx=ast.Store()) + named = ast.NamedExpr(target=target, value=ensure_ast(state.stack[-1])) + return replace(state, stack=state.stack[:-2] + [named]) + + +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_312) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_313) +@register_handler("STORE_GLOBAL", version=PythonVersion.PY_314) +def handle_store_global( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A comprehension has no globals of its own, so the only way it stores one + # is an assignment expression, which binds in the enclosing scope. + assert _is_assignment_expression(state), ( + "STORE_GLOBAL outside an assignment expression" + ) + return _handle_assignment_expression(state, instr) + + +@register_handler("STORE_DEREF", version=PythonVersion.PY_312) +@register_handler("STORE_DEREF", version=PythonVersion.PY_313) +@register_handler("STORE_DEREF", version=PythonVersion.PY_314) +def handle_store_deref( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_DEREF stores into a closure variable: either an assignment + # expression binding in an enclosing function, or a loop target that an + # inner comprehension captures. + if _is_assignment_expression(state): + return _handle_assignment_expression(state, instr) + return handle_store_fast(state, instr) + + +@register_handler("STORE_FAST", version=PythonVersion.PY_312) +@register_handler("STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + if _is_assignment_expression(state): + # An assignment expression whose target is local to this code object, + # as in a comprehension inlined into the lambda that binds the name. + return _handle_assignment_expression(state, instr) + + if isinstance(state.stack[-1], ast.Name) and state.stack[-1].id == instr.argval: + # If the variable is already on the stack, we can skip adding it again + # This is common in nested comprehensions where the same variable is reused + return replace(state, stack=state.stack[:-1]) + + if not isinstance(state.stack[-1], TargetHole): + return _bind_local(state, instr.argval, state.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=instr.argval, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1]) + + +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_store_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_LOAD_FAST stores and then loads the same variable (optimization) + # The instruction has two names: store_name and load_name + # In Python 3.13, this is often used for loop variables + + # In Python 3.13, the instruction argument contains both names + # argval should be a tuple (store_name, load_name) + assert isinstance(instr.argval, tuple) + store_name, load_name = instr.argval + + if _is_assignment_expression(state): + # `(name := value)` whose result is read straight back, as in + # `(z := w) + z`. The duplicate becomes the assignment expression and + # the reload becomes a plain reference to the name just bound. + named = ast.NamedExpr( + target=ast.Name(id=store_name, ctx=ast.Store()), + value=ensure_ast(state.stack[-1]), + ) + reload = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=state.stack[:-2] + [named, reload]) + + if not isinstance(state.stack[-1], TargetHole): + # A plain assignment followed by a load, as 3.14 emits when it unrolls + # a single-iteration loop over a literal. + bound = _bind_local(state, store_name, state.stack[-1]) + return replace(bound, stack=bound.stack + [_read_local(bound, load_name)]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=store_name, ctx=ast.Store()) + ) + new_var = ast.Name(id=load_name, ctx=ast.Load()) + return replace(state, stack=new_stack[:-1] + [new_var]) + + +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_313) +@register_handler("STORE_FAST_STORE_FAST", version=PythonVersion.PY_314) +def handle_store_fast_store_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # STORE_FAST_STORE_FAST stores STACK[-1] into the first named variable and + # STACK[-2] into the second. It is emitted for unpacking targets, so both + # values are loop-target holes belonging to the same comprehension. + assert isinstance(instr.argval, tuple) + first_name, second_name = instr.argval + + if not isinstance(state.stack[-1], TargetHole): + # Not loop targets: a pair of plain assignments, as 3.14 emits when it + # unrolls a single-iteration loop over a literal. + bound = _bind_local(state, first_name, state.stack[-1]) + return _bind_local(bound, second_name, bound.stack[-1]) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Name(id=first_name, ctx=ast.Store()) + ) + new_stack = _bind_target_hole( + new_stack, new_stack[-2], ast.Name(id=second_name, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-2]) + + +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_312) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_AND_CLEAR", version=PythonVersion.PY_314) +def handle_load_fast_and_clear( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_AND_CLEAR pushes a local variable onto the stack and clears it + # For AST reconstruction, we treat this the same as LOAD_FAST + return replace(state, stack=state.stack + [_read_local(state, instr.argval)]) + + +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_313) +@register_handler("LOAD_FAST_LOAD_FAST", version=PythonVersion.PY_314) +def handle_load_fast_load_fast( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_FAST_LOAD_FAST loads two variables (optimization in Python 3.13) + # The instruction argument contains both variable names + if isinstance(instr.argval, tuple): + var1, var2 = instr.argval + else: + # Fallback: assume both names are the same + var1 = var2 = instr.argval + + new_stack = state.stack + [_read_local(state, var1), _read_local(state, var2)] + + return replace(state, stack=new_stack) + + +@register_handler("MAKE_CELL", version=PythonVersion.PY_312) +@register_handler("MAKE_CELL", version=PythonVersion.PY_313) +@register_handler("MAKE_CELL", version=PythonVersion.PY_314) +def handle_make_cell( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_CELL creates a new cell in slot i for closure variables + # This is used when variables from outer scopes are captured by inner scopes + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track low-level closure details + return state + + +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_312) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_313) +@register_handler("COPY_FREE_VARS", version=PythonVersion.PY_314) +def handle_copy_free_vars( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY_FREE_VARS copies n free (closure) variables from the closure into the frame + # This removes the need for special code on the caller's side when calling closures + # For AST reconstruction purposes, this is just a variable scoping mechanism + # that we can ignore since the AST doesn't track runtime variable management + return state + + +# ============================================================================ +# STACK MANAGEMENT HANDLERS +# ============================================================================ + + +@register_handler("POP_TOP", version=PythonVersion.PY_312) +@register_handler("POP_TOP", version=PythonVersion.PY_313) +@register_handler("POP_TOP", version=PythonVersion.PY_314) +def handle_pop_top( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_TOP removes the top item from the stack + # In generators, often used after YIELD_VALUE + # Also used to clean up the duplicated middle value in failed chained comparisons + new_stack = state.stack[:-1] + return replace(state, stack=new_stack) + + +# Python 3.13 replacement for stack manipulation +@register_handler("SWAP", version=PythonVersion.PY_312) +@register_handler("SWAP", version=PythonVersion.PY_313) +@register_handler("SWAP", version=PythonVersion.PY_314) +def handle_swap( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SWAP exchanges the top two stack items (replaces ROT_TWO in many cases) + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + + if depth > stack_size: + # Not enough items on stack - this might be a pattern where some items were optimized away + # For AST reconstruction, we can often ignore certain stack manipulations + return state + + # For other depths, swap TOS with the item at specified depth + assert depth <= stack_size, f"SWAP depth {depth} exceeds stack size {stack_size}" + idx = stack_size - depth + new_stack = state.stack.copy() + new_stack[-1], new_stack[idx] = new_stack[idx], new_stack[-1] + return replace(state, stack=new_stack) + + +@register_handler("COPY", version=PythonVersion.PY_312) +@register_handler("COPY", version=PythonVersion.PY_313) +@register_handler("COPY", version=PythonVersion.PY_314) +def handle_copy( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # COPY duplicates the item at the specified depth + assert instr.arg is not None + depth = instr.arg + stack_size = len(state.stack) + if depth > stack_size: + raise ValueError(f"COPY depth {depth} exceeds stack size {stack_size}") + idx = stack_size - depth + copied_item = state.stack[idx] + new_stack = state.stack + [copied_item] + return replace(state, stack=new_stack) + + +@register_handler("PUSH_NULL", version=PythonVersion.PY_312) +@register_handler("PUSH_NULL", version=PythonVersion.PY_313) +@register_handler("PUSH_NULL", version=PythonVersion.PY_314) +def handle_push_null( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return replace(state, stack=state.stack + [Null()]) + + +# ============================================================================ +# BINARY ARITHMETIC/LOGIC OPERATION HANDLERS +# ============================================================================ + + +def handle_binop( + op: ast.operator, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + [ast.BinOp(left=left, op=op, right=right)] + return replace(state, stack=new_stack) + + +# Python 3.12+ BINARY_OP handler +@register_handler("BINARY_OP", version=PythonVersion.PY_312) +@register_handler("BINARY_OP", version=PythonVersion.PY_313) +def handle_binary_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_OP in Python 3.12+ consolidates all binary operations + # The operation type is determined by the instruction argument + assert instr.arg is not None + + # Map argument values to AST operators based on Python 3.12+ implementation + op_map: collections.abc.Mapping[int, ast.operator] = { + 0: ast.Add(), # + + 1: ast.BitAnd(), # & + 2: ast.FloorDiv(), # // + 3: ast.LShift(), # << + 4: ast.MatMult(), # @ + 5: ast.Mult(), # * + 6: ast.Mod(), # % + 7: ast.BitOr(), # | + 8: ast.Pow(), # ** + 9: ast.RShift(), # >> + 10: ast.Sub(), # - + 11: ast.Div(), # / + 12: ast.BitXor(), # ^ + } + + op = op_map.get(instr.arg) + if op is None: + raise TypeError(f"Unknown binary operation: {instr.arg}") + + return handle_binop(op, state, instr) + + +# 3.14 folded subscripting into BINARY_OP; `dis._nb_ops` names the oparg +# NB_SUBSCR. Looked up rather than hard-coded, since it sits past the in-place +# operators and so moves whenever one is added. +_NB_OPS: list[tuple[str, str]] = getattr(dis, "_nb_ops", []) +NB_SUBSCR: int | None = next( + (i for i, (name, _) in enumerate(_NB_OPS) if name == "NB_SUBSCR"), None +) + + +@register_handler("BINARY_OP", version=PythonVersion.PY_314) +def handle_binary_op_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # As in 3.13, except that BINARY_OP now also implements `a[b]`, which used + # to be its own BINARY_SUBSCR opcode. + if instr.arg is not None and instr.arg == NB_SUBSCR: + return handle_binary_subscr(state, instr) + return handle_binary_op(state, instr) + + +@register_handler("LOAD_SMALL_INT", version=PythonVersion.PY_314) +def handle_load_small_int( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_SMALL_INT pushes an int in range(256) held in the oparg itself, + # rather than going through co_consts. + assert isinstance(instr.argval, int) + return replace(state, stack=state.stack + [ensure_ast(instr.argval)]) + + +@register_handler("LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A borrowed reference differs only in ownership, which the AST does not + # model, so this is LOAD_FAST as far as reconstruction is concerned. + return handle_load_fast(state, instr) + + +@register_handler("LOAD_FAST_BORROW_LOAD_FAST_BORROW", version=PythonVersion.PY_314) +def handle_load_fast_borrow_load_fast_borrow( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + return handle_load_fast_load_fast(state, instr) + + +@register_handler("LOAD_COMMON_CONSTANT", version=PythonVersion.PY_314) +def handle_load_common_constant( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Pushes one of a small hardcoded set of constants. In a comprehension this + # only shows up in the guard of an inlined builtin; see CommonConstant. + name = getattr(instr.argval, "__name__", str(instr.argval)) + return replace(state, stack=state.stack + [CommonConstant(id=name)]) + + +@register_handler("NOT_TAKEN", version=PythonVersion.PY_314) +def handle_not_taken( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # A no-op marking the not-taken edge of a branch for sys.monitoring. + return state + + +@register_handler("POP_ITER", version=PythonVersion.PY_314) +def handle_pop_iter( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_ITER discards the exhausted iterator that FOR_ITER left behind. In + # 3.13 the same cleanup was spelled END_FOR followed by POP_TOP. + return replace(state, stack=state.stack[:-1]) + + +# ============================================================================ +# UNARY OPERATION HANDLERS +# ============================================================================ + + +def handle_unary_op( + op: ast.unaryop, state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + operand = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + [ast.UnaryOp(op=op, operand=operand)] + return replace(state, stack=new_stack) + + +UNARY_OPS: dict[str, ast.unaryop] = { + "UNARY_NEGATIVE": ast.USub(), + "UNARY_INVERT": ast.Invert(), + "UNARY_NOT": ast.Not(), +} + +# These three behave identically in every dialect this module supports; 3.13's +# "requires an exact bool operand" note on UNARY_NOT constrains the operand, not +# the reconstruction. +for _opname, _op in UNARY_OPS.items(): + for _version in ( + PythonVersion.PY_312, + PythonVersion.PY_313, + PythonVersion.PY_314, + ): + register_handler( + _opname, functools.partial(handle_unary_op, _op), version=_version + ) + + +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_313) +@register_handler("CONVERT_VALUE", version=PythonVersion.PY_314) +def handle_convert_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CONVERT_VALUE applies a conversion to the value on top of stack + # Used for f-string conversions like !r, !s, !a + # The conversion type is stored in instr.arg: + # 0 = None, 1 = str (!s), 2 = repr (!r), 3 = ascii (!a) + assert len(state.stack) > 0, "CONVERT_VALUE requires a value on stack" + assert instr.arg is not None, "CONVERT_VALUE requires conversion type" + + # Wrap the value with conversion information + value = state.stack[-1] + converted = ConvertedValue(value, instr.arg) + new_stack = state.stack[:-1] + [converted] + + return replace(state, stack=new_stack) + + +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_312) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_313) +@register_handler("CALL_INTRINSIC_1", version=PythonVersion.PY_314) +def handle_call_intrinsic_1( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_INTRINSIC_1 calls an intrinsic function with one argument + if instr.argrepr == "INTRINSIC_LIST_TO_TUPLE": + assert isinstance(state.stack[-1], ast.List), ( + "Expected a list for LIST_TO_TUPLE" + ) + tuple_node = ast.Tuple(elts=state.stack[-1].elts, ctx=ast.Load()) + return replace(state, stack=state.stack[:-1] + [tuple_node]) + elif instr.argrepr == "INTRINSIC_UNARY_POSITIVE": + assert len(state.stack) > 0 + new_val = ast.UnaryOp(op=ast.UAdd(), operand=state.stack[-1]) + return replace(state, stack=state.stack[:-1] + [new_val]) + elif instr.argrepr == "INTRINSIC_STOPITERATION_ERROR": + return state + else: + raise TypeError(f"Unsupported generator intrinsic operation: {instr.argrepr}") + + +@register_handler("TO_BOOL", version=PythonVersion.PY_313) +@register_handler("TO_BOOL", version=PythonVersion.PY_314) +def handle_to_bool( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # TO_BOOL converts the top stack item to a boolean + # For AST reconstruction, we typically don't need an explicit bool() call + # since the boolean context is usually handled by the conditional jump that follows + # However, for some cases we might need to preserve the explicit conversion + + # For now, leave the value as-is since the jump instruction will handle the boolean logic + return state + + +# ============================================================================ +# COMPARISON OPERATION HANDLERS +# ============================================================================ + +CMP_OPMAP: dict[str, ast.cmpop] = { + "<": ast.Lt(), + "<=": ast.LtE(), + ">": ast.Gt(), + ">=": ast.GtE(), + "==": ast.Eq(), + "!=": ast.NotEq(), +} + + +@register_handler("COMPARE_OP", version=PythonVersion.PY_312) +@register_handler("COMPARE_OP", version=PythonVersion.PY_313) +@register_handler("COMPARE_OP", version=PythonVersion.PY_314) +def handle_compare_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None and instr.argval in dis.cmp_op, ( + f"Unsupported comparison operation: {instr.argval}" + ) + + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # Map comparison operation codes to AST operators + op_name = instr.argval + compare_node = ast.Compare(left=left, ops=[CMP_OPMAP[op_name]], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("CONTAINS_OP", version=PythonVersion.PY_312) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_313) +@register_handler("CONTAINS_OP", version=PythonVersion.PY_314) +def handle_contains_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) # Container + left = ensure_ast(state.stack[-2]) # Item to check + + # instr.arg determines if it's 'in' (0) or 'not in' (1) + op = ast.NotIn() if instr.arg else ast.In() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +@register_handler("IS_OP", version=PythonVersion.PY_312) +@register_handler("IS_OP", version=PythonVersion.PY_313) +@register_handler("IS_OP", version=PythonVersion.PY_314) +def handle_is_op( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + right = ensure_ast(state.stack[-1]) + left = ensure_ast(state.stack[-2]) + + # instr.arg determines if it's 'is' (0) or 'is not' (1) + op = ast.IsNot() if instr.arg else ast.Is() + + compare_node = ast.Compare(left=left, ops=[op], comparators=[right]) + new_stack = state.stack[:-2] + [compare_node] + return replace(state, stack=new_stack) + + +# ============================================================================ +# FUNCTION CALL HANDLERS +# ============================================================================ + + +@register_handler("KW_NAMES", version=PythonVersion.PY_312) +def handle_kw_names( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # KW_NAMES names the trailing arguments of the CALL that follows it. + # Python 3.13 replaced this pair with a single CALL_KW instruction. + assert isinstance(instr.argval, tuple), "KW_NAMES requires a tuple of names" + assert all(isinstance(name, str) for name in instr.argval) + assert state.kw_names is None, "KW_NAMES must be consumed by the following CALL" + return replace(state, kw_names=instr.argval) + + +@register_handler("CALL", version=PythonVersion.PY_312) +def handle_call_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL in Python 3.12 handles both function and method calls + # Stack layout: [..., callable or self, callable or NULL] + assert instr.arg is not None + arg_count: int = instr.arg + + # Check if this is a method call (no NULL on top) + if isinstance(state.stack[-arg_count - 2], Null): + # Regular function call: [..., NULL, callable, *args] + func = ensure_ast(state.stack[-arg_count - 1]) + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + new_stack = state.stack[: -arg_count - 2] + else: + # Method call: [..., callable, self, *args] + func = ensure_ast(state.stack[-arg_count - 2]) + self_arg = ensure_ast(state.stack[-arg_count - 1]) + remaining_args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] + if arg_count > 0 + else [] + ) + args = [self_arg] + remaining_args + new_stack = state.stack[: -arg_count - 2] + + # A preceding KW_NAMES names the trailing `len(kw_names)` positional slots. + keywords: list[ast.keyword] = [] + if state.kw_names is not None: + assert 0 < len(state.kw_names) <= arg_count + keywords = [ + ast.keyword(arg=name, value=value) + for name, value in zip(state.kw_names, args[-len(state.kw_names) :]) + ] + args = args[: -len(state.kw_names)] + + if isinstance(func, CompLambda): + assert len(args) == 1 and not keywords + return replace(state, stack=new_stack + [func.inline(args[0])], kw_names=None) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack, kw_names=None) + + +@register_handler("CALL", version=PythonVersion.PY_313) +@register_handler("CALL", version=PythonVersion.PY_314) +def handle_call( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL pops function and arguments from stack (replaces CALL_FUNCTION in Python 3.13) + assert instr.arg is not None + arg_count: int = instr.arg + + func = ensure_ast(state.stack[-arg_count - 2]) + + # Pop arguments and function + args = ( + [ensure_ast(arg) for arg in state.stack[-arg_count:]] if arg_count > 0 else [] + ) + if not isinstance(state.stack[-arg_count - 1], Null): + args = [ensure_ast(state.stack[-arg_count - 1])] + args + + new_stack = state.stack[: -arg_count - 2] + if isinstance(func, CompLambda): + assert len(args) == 1 + return replace(state, stack=new_stack + [func.inline(args[0])]) + else: + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=[]) + new_stack = new_stack + [call_node] + return replace(state, stack=new_stack) + + +@register_handler("CALL_KW", version=PythonVersion.PY_313) +@register_handler("CALL_KW", version=PythonVersion.PY_314) +def handle_call_kw( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # CALL_KW pops function, arguments, and keyword names from stack + assert instr.arg is not None + arg_count: int = instr.arg + assert arg_count > 0, "CALL_KW requires at least one argument" + + func = ensure_ast(state.stack[-arg_count - 3]) + assert not isinstance(func, CompLambda | Null) + + kw_names = state.stack[-1] + assert isinstance(kw_names, ast.Tuple), "Expected a tuple of keyword names" + assert len(kw_names.elts) > 0, "Expected at least one keyword name" + + # Pop arguments, function, and keyword names + keywords = [] + for i, kw in enumerate(reversed(kw_names.elts)): + assert isinstance(kw, ast.Constant) and isinstance(kw.value, str) + keywords += [ast.keyword(arg=kw.value, value=ensure_ast(state.stack[-2 - i]))] + keywords.reverse() + + args = [ensure_ast(a) for a in state.stack[-arg_count - 1 : -len(keywords) - 1]] + if not isinstance(state.stack[-arg_count - 2], Null): + args = [ensure_ast(state.stack[-arg_count - 2])] + args + + # Create function call AST + call_node = ast.Call(func=func, args=args, keywords=keywords) + new_stack = state.stack[: -arg_count - 3] + [call_node] + return replace(state, stack=new_stack) + + +# Flags shared by MAKE_FUNCTION (3.12) and SET_FUNCTION_ATTRIBUTE (3.13) +MAKE_FUNCTION_DEFAULTS = 0x01 +MAKE_FUNCTION_KWDEFAULTS = 0x02 +MAKE_FUNCTION_ANNOTATIONS = 0x04 +MAKE_FUNCTION_CLOSURE = 0x08 +MAKE_FUNCTION_ANNOTATE = 0x10 # added in 3.14 +MAKE_FUNCTION_FLAGS = ( + MAKE_FUNCTION_DEFAULTS, + MAKE_FUNCTION_KWDEFAULTS, + MAKE_FUNCTION_ANNOTATIONS, + MAKE_FUNCTION_CLOSURE, +) + + +def _apply_function_attribute( + func: ast.Lambda | CompLambda, flag: int, value: ast.expr +) -> ast.Lambda | CompLambda: + """Attach one function attribute to a reconstructed lambda.""" + if flag == MAKE_FUNCTION_CLOSURE: + # Free variables are already spelled by name in the reconstructed body. + return func + if flag == MAKE_FUNCTION_ANNOTATE: + # A lambda has no annotations, and the AST does not carry the lazy + # annotate function 3.14 attaches to annotated functions. + return func + + assert isinstance(func, ast.Lambda) and not isinstance(func, CompLambda), ( + "Only lambdas carry defaults; comprehensions take exactly one argument" + ) + + if flag == MAKE_FUNCTION_DEFAULTS: + # A tuple of defaults for the *trailing* positional parameters. + assert isinstance(value, ast.Tuple), "Expected a tuple of default values" + func.args.defaults = list(value.elts) + elif flag == MAKE_FUNCTION_KWDEFAULTS: + # A dict mapping keyword-only parameter names to their defaults. + assert isinstance(value, ast.Dict), "Expected a dict of keyword defaults" + by_name = { + key.value: val + for key, val in zip(value.keys, value.values) + if isinstance(key, ast.Constant) + } + func.args.kw_defaults = [by_name.get(a.arg) for a in func.args.kwonlyargs] + else: + raise NotImplementedError("Function annotations are not supported") + + return func + + +def _split_callable( + first: ast.expr, second: ast.expr +) -> tuple[ast.expr, ast.expr | None]: + """Separate the callable from the NULL-or-self slot beside it. + + Which of the two comes first varies: LOAD_GLOBAL and LOAD_ATTR report the + order in their argrepr, and it differs between 3.13 and 3.14. + """ + if isinstance(first, Null): + return second, None + elif isinstance(second, Null): + return first, None + else: + return first, second + + +def _build_variadic_call( + func: ast.expr, + self_arg: ast.expr | None, + positional: ast.expr, + keyword_mapping: ast.expr | None, +) -> ast.Call: + """Assemble `func(*positional, **keyword_mapping)`. + + CALL_FUNCTION_EX receives its arguments already collected into a sequence + and a mapping, with the original mix of plain and starred arguments no + longer distinguishable. Spelling every argument as unpacked reproduces the + call exactly, even where the source did not use `*` for all of them. + """ + args: list[ast.expr] = [] if self_arg is None else [ensure_ast(self_arg)] + args.append(ast.Starred(value=ensure_ast(positional), ctx=ast.Load())) + keywords = ( + [] + if keyword_mapping is None + else [ast.keyword(arg=None, value=ensure_ast(keyword_mapping))] + ) + return ast.Call(func=ensure_ast(func), args=args, keywords=keywords) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_312) +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_313) +def handle_call_function_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # Stack: callable and NULL-or-self, the positional sequence, and -- only + # when the low bit of the oparg is set -- the keyword mapping. + size = 4 if instr.arg else 3 + keyword_mapping = state.stack[-1] if instr.arg else None + positional = state.stack[-2] if instr.arg else state.stack[-1] + func, self_arg = _split_callable(state.stack[-size], state.stack[-size + 1]) + + call = _build_variadic_call(func, self_arg, positional, keyword_mapping) + return replace(state, stack=state.stack[:-size] + [call]) + + +@register_handler("CALL_FUNCTION_EX", version=PythonVersion.PY_314) +def handle_call_function_ex_314( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # 3.14 always reserves the keyword-mapping slot, pushing NULL into it when + # the call has no `**` argument, so the layout is a fixed four slots. + keyword_mapping = None if isinstance(state.stack[-1], Null) else state.stack[-1] + func, self_arg = _split_callable(state.stack[-4], state.stack[-3]) + + call = _build_variadic_call(func, self_arg, state.stack[-2], keyword_mapping) + return replace(state, stack=state.stack[:-4] + [call]) + + +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_312) +def handle_make_function_312( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.12 uses flags to determine stack consumption. + # Unlike 3.10 there is no qualified name on the stack, and unlike 3.13 the + # extra attributes travel with this instruction rather than with a following + # SET_FUNCTION_ATTRIBUTE. They are pushed in ascending flag order, below the + # code object. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + set_flags = [flag for flag in MAKE_FUNCTION_FLAGS if instr.arg & flag] + attributes = state.stack[-1 - len(set_flags) : -1] + + func = copy.deepcopy(state.stack[-1]) + for flag, value in zip(set_flags, attributes): + func = _apply_function_attribute(func, flag, value) + + new_stack = state.stack[: -1 - len(set_flags)] + [func] + return replace(state, stack=new_stack) + + +# Python 3.13 version +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_313) +@register_handler("MAKE_FUNCTION", version=PythonVersion.PY_314) +def handle_make_function( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # MAKE_FUNCTION in Python 3.13 is simplified: it only takes a code object from the stack + # and creates a function from it. No flags, no extra attributes on the stack. + # All extra attributes are handled by separate SET_FUNCTION_ATTRIBUTE instructions. + + # Pop the function object from the stack (it's the only thing expected) + # Conversion from CodeType to ast.Lambda should have happened already + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + return state + + +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_313) +@register_handler("SET_FUNCTION_ATTRIBUTE", version=PythonVersion.PY_314) +def handle_set_function_attribute( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # SET_FUNCTION_ATTRIBUTE sets one attribute on a function object. Python + # 3.13 uses it in place of the MAKE_FUNCTION flags; the stack holds the + # attribute value below the function, and only the function is left behind. + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Lambda | CompLambda), ( + "Expected a function object (Lambda or CompLambda) on the stack." + ) + + func = _apply_function_attribute( + copy.deepcopy(state.stack[-1]), instr.arg, state.stack[-2] + ) + return replace(state, stack=state.stack[:-2] + [func]) + + +# ============================================================================ +# OBJECT ACCESS HANDLERS +# ============================================================================ + + +@register_handler("LOAD_ATTR", version=PythonVersion.PY_312) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_313) +@register_handler("LOAD_ATTR", version=PythonVersion.PY_314) +def handle_load_attr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LOAD_ATTR loads an attribute from the object on top of stack + obj = ensure_ast(state.stack[-1]) + attr_name = instr.argval + + # Create attribute access AST + attr_node = ast.Attribute(value=obj, attr=attr_name, ctx=ast.Load()) + if instr.argrepr.endswith(" + NULL|self"): + new_stack = state.stack[:-1] + [attr_node, Null()] + elif instr.argrepr.startswith("NULL|self + "): + new_stack = state.stack[:-1] + [Null(), attr_node] + else: + new_stack = state.stack[:-1] + [attr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_312) +@register_handler("BINARY_SUBSCR", version=PythonVersion.PY_313) +def handle_binary_subscr( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SUBSCR implements obj[index] - pops index and obj from stack + index = ensure_ast(state.stack[-1]) # Index is on top + obj = ensure_ast(state.stack[-2]) # Object is below index + new_stack = state.stack[:-2] + + # Create subscript access AST + subscr_node = ast.Subscript(value=obj, slice=index, ctx=ast.Load()) + new_stack = new_stack + [subscr_node] + return replace(state, stack=new_stack) + + +@register_handler("BINARY_SLICE", version=PythonVersion.PY_312) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_313) +@register_handler("BINARY_SLICE", version=PythonVersion.PY_314) +def handle_binary_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BINARY_SLICE implements obj[start:end] - pops start, end, and obj from stack + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + container = ensure_ast(state.stack[-3]) # Object is below start and end + sliced = ast.Subscript( + value=container, + slice=ast.Slice(lower=start, upper=end, step=None), + ctx=ast.Load(), + ) + new_stack = state.stack[:-3] + [sliced] + return replace(state, stack=new_stack) + + +# ============================================================================ +# OTHER CONTAINER BUILDING HANDLERS +# ============================================================================ + + +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_312) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_313) +@register_handler("UNPACK_SEQUENCE", version=PythonVersion.PY_314) +def handle_unpack_sequence( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_SEQUENCE splits a comprehension loop target into `arg` sub-targets, + # as in ((k, v) for k, v in items). The names are not known yet, so the + # single target hole is refined into a tuple of fresh holes, which the + # following STORE_* instructions bind one at a time. + # + # CPython pushes the unpacked values right-to-left, so element 0 ends up on + # top of the stack and is consumed by the first STORE_*. + assert instr.arg is not None + unpack_count: int = instr.arg + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value rather than a loop target, as 3.14 emits + # when it unrolls a single-iteration loop over a literal. + elements = _literal_elements(state.stack[-1], unpack_count) + return replace(state, stack=state.stack[:-1] + list(reversed(elements))) + + holes = [TargetHole() for _ in range(unpack_count)] + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=list(holes), ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("UNPACK_EX", version=PythonVersion.PY_312) +@register_handler("UNPACK_EX", version=PythonVersion.PY_313) +@register_handler("UNPACK_EX", version=PythonVersion.PY_314) +def handle_unpack_ex( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # UNPACK_EX handles a starred target, as in ((a, b) for a, *b in pairs). + # The low byte of the argument counts the targets before the starred one and + # the high byte counts those after it; the starred target itself collects + # whatever is left over. As with UNPACK_SEQUENCE the values are pushed + # right-to-left, so the first target ends up on top of the stack. + assert instr.arg is not None + before, after = instr.arg & 0xFF, instr.arg >> 8 + + if not isinstance(state.stack[-1], TargetHole): + # Destructuring a known value; the starred target collects the middle. + elements = _literal_elements(state.stack[-1]) + assert len(elements) >= before + after, "Too few values to unpack" + middle = elements[before : len(elements) - after] + unpacked: list[ast.expr] = [ + *elements[:before], + ast.List(elts=list(middle), ctx=ast.Load()), + *elements[len(elements) - after :], + ] + return replace(state, stack=state.stack[:-1] + list(reversed(unpacked))) + + holes = [TargetHole() for _ in range(before + 1 + after)] + elts: list[ast.expr] = list(holes) + elts[before] = ast.Starred(value=holes[before], ctx=ast.Store()) + + new_stack = _bind_target_hole( + state.stack, state.stack[-1], ast.Tuple(elts=elts, ctx=ast.Store()) + ) + return replace(state, stack=new_stack[:-1] + list(reversed(holes))) + + +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_312) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_313) +@register_handler("BUILD_TUPLE", version=PythonVersion.PY_314) +def handle_build_tuple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + assert instr.arg is not None + tuple_size: int = instr.arg + # Pop elements for the tuple + elements = ( + [ensure_ast(elem) for elem in state.stack[-tuple_size:]] + if tuple_size > 0 + else [] + ) + new_stack = state.stack[:-tuple_size] if tuple_size > 0 else state.stack + + # Create tuple AST + tuple_node = ast.Tuple(elts=elements, ctx=ast.Load()) + new_stack = new_stack + [tuple_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_SLICE", version=PythonVersion.PY_312) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_313) +@register_handler("BUILD_SLICE", version=PythonVersion.PY_314) +def handle_build_slice( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_SLICE creates a slice object from the top of the stack + # The number of elements to pop is determined by the instruction argument + assert instr.arg is not None + slice_size: int = instr.arg + + if slice_size == 2: + # Slice with start and end: [start, end] + end = ensure_ast(state.stack[-1]) + start = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + slice_node = ast.Slice(lower=start, upper=end, step=None) + elif slice_size == 3: + # Slice with start, end, and step: [start, end, step] + step = ensure_ast(state.stack[-1]) + end = ensure_ast(state.stack[-2]) + start = ensure_ast(state.stack[-3]) + new_stack = state.stack[:-3] + slice_node = ast.Slice(lower=start, upper=end, step=step) + else: + raise ValueError(f"Unsupported slice size: {slice_size}") + + # Create slice AST + new_stack = new_stack + [slice_node] + return replace(state, stack=new_stack) + + +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_312) +@register_handler("BUILD_CONST_KEY_MAP", version=PythonVersion.PY_313) +def handle_build_const_key_map( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_CONST_KEY_MAP builds a dictionary with constant keys + # The keys are in a tuple on TOS, values are on the stack below + assert instr.arg is not None + assert isinstance(state.stack[-1], ast.Tuple), "Expected a tuple of keys" + map_size: int = instr.arg + # Pop the keys tuple and values + keys_tuple: ast.Tuple = state.stack[-1] + keys: list[ast.expr | None] = [ensure_ast(key) for key in keys_tuple.elts] + values = [ensure_ast(val) for val in state.stack[-map_size - 1 : -1]] + new_stack = state.stack[: -map_size - 1] + + # Create dictionary AST + dict_node = ast.Dict(keys=keys, values=values) + new_stack = new_stack + [dict_node] + return replace(state, stack=new_stack) + + +@register_handler("LIST_EXTEND", version=PythonVersion.PY_312) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_313) +@register_handler("LIST_EXTEND", version=PythonVersion.PY_314) +def handle_list_extend( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # LIST_EXTEND appends the contents of the iterable at TOS to the list + # further down the stack. That list is either the empty ListComp that + # BUILD_LIST(0) optimistically created -- a list display, not a + # comprehension after all -- or a partly built argument list for a call + # with a starred argument. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # A literal iterable contributes its elements directly; anything else has to + # stay unpacked, as in `[*whatever]`. + elements: list[ast.expr] + if isinstance(update, ast.Tuple | ast.List): + elements = [ensure_ast(e) for e in update.elts] + else: + elements = [ast.Starred(value=ensure_ast(update), ctx=ast.Load())] + + if isinstance(target, ast.ListComp) and not target.generators: + merged = ast.List(elts=elements, ctx=ast.Load()) + else: + assert isinstance(target, ast.List), "LIST_EXTEND expects a list to extend" + merged = ast.List(elts=list(target.elts) + elements, ctx=ast.Load()) + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = merged + return replace(state, stack=new_stack) + + +@register_handler("DICT_MERGE", version=PythonVersion.PY_312) +@register_handler("DICT_MERGE", version=PythonVersion.PY_313) +@register_handler("DICT_MERGE", version=PythonVersion.PY_314) +def handle_dict_merge( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # DICT_MERGE folds the mapping at TOS into the one below it, rejecting + # duplicate keys. It assembles the keyword arguments of a call using `**`. + update = state.stack[-1] + target = state.stack[-instr.argval - 1] + + # An `ast.Dict` entry with a key of None is `**value`, which is how a + # mapping that is not a literal has to be spliced in. + def entries(node: ast.expr) -> tuple[list[ast.expr | None], list[ast.expr]]: + if isinstance(node, ast.Dict): + return list(node.keys), list(node.values) + return [None], [ensure_ast(node)] + + if isinstance(target, ast.DictComp) and not target.generators: + # BUILD_MAP(0) guessed at a dict comprehension; it was a `**` argument. + keys, values = entries(update) + else: + assert isinstance(target, ast.Dict), "DICT_MERGE expects a dict to merge into" + target_keys, target_values = entries(target) + update_keys, update_values = entries(update) + keys, values = target_keys + update_keys, target_values + update_values + + new_stack = state.stack[:-1] + new_stack[-instr.argval] = ast.Dict(keys=keys, values=values) + return replace(state, stack=new_stack) + + +@register_handler("SET_UPDATE", version=PythonVersion.PY_312) +@register_handler("SET_UPDATE", version=PythonVersion.PY_313) +@register_handler("SET_UPDATE", version=PythonVersion.PY_314) +def handle_set_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The set being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_SET, + # while the actual result expression is in the stack where the set "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.SetComp) + assert isinstance(state.stack[-1], ast.Tuple | ast.List | ast.Set) + + new_val = ast.Set(elts=[ensure_ast(e) for e in state.stack[-1].elts]) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("DICT_UPDATE", version=PythonVersion.PY_312) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_313) +@register_handler("DICT_UPDATE", version=PythonVersion.PY_314) +def handle_dict_update( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # The dict being extended is actually in state.result instead of the stack + # because it was initially recognized as a list comprehension in BUILD_MAP, + # while the actual result expression is in the stack where the dict "should be" + # and needs to be put back into the state result slot + assert isinstance(state.stack[-instr.argval - 1], ast.DictComp) + assert isinstance(state.stack[-1], ast.Dict) + + new_val = ast.Dict( + keys=[ensure_ast(e) for e in state.stack[-1].keys], + values=[ensure_ast(e) for e in state.stack[-1].values], + ) + new_stack = state.stack[:-2] + [new_val] + + return replace(state, stack=new_stack) + + +@register_handler("BUILD_STRING", version=PythonVersion.PY_312) +@register_handler("BUILD_STRING", version=PythonVersion.PY_313) +@register_handler("BUILD_STRING", version=PythonVersion.PY_314) +def handle_build_string( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # BUILD_STRING concatenates strings from the stack + # For f-strings, it combines FormattedValue and Constant nodes + assert instr.arg is not None + string_size: int = instr.arg + + if string_size == 0: + # Empty string case + new_stack = state.stack + [ast.Constant(value="")] + return replace(state, stack=new_stack) + + # Pop elements for the string + elements = [ensure_ast(elem) for elem in state.stack[-string_size:]] + new_stack = state.stack[:-string_size] + + # Check if this is an f-string build (has FormattedValue nodes) + # or a regular string concatenation + if any(isinstance(elem, ast.JoinedStr) for elem in elements): + # This is an f-string - create JoinedStr + values = [] + for elem in elements: + if isinstance(elem, ast.JoinedStr): + values.extend(elem.values) + else: + values.append(elem) + return replace(state, stack=new_stack + [ast.JoinedStr(values=values)]) + elif all(isinstance(elem, ast.Constant) for elem in elements): + # This is regular string concatenation or format spec building + # If all elements are constants, we might be building a format spec + # Concatenate the constant strings + assert all( + isinstance(elem, ast.Constant) and isinstance(elem.value, str) + for elem in elements + ) + concat_str = "".join( + elem.value + for elem in elements + if isinstance(elem, ast.Constant) and isinstance(elem.value, str) + ) + return replace(state, stack=new_stack + [ast.Constant(value=concat_str)]) + else: + raise TypeError("Should not be here?") + + +@register_handler("FORMAT_VALUE", version=PythonVersion.PY_312) +def handle_format_value( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_VALUE formats a string with a value in Python 3.12 + # Flag bits: (flags & 0x03) = conversion, (flags & 0x04) = has format spec + assert instr.arg is not None, "FORMAT_VALUE requires flags argument" + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_VALUE" + + flags = instr.arg + + # Check if there's a format specification + has_format_spec = bool(flags & 0x04) + + if has_format_spec: + # Pop format spec and value + assert len(state.stack) >= 2, ( + "FORMAT_VALUE with format spec needs 2 stack items" + ) + format_spec = ensure_ast(state.stack[-1]) + value = ensure_ast(state.stack[-2]) + new_stack = state.stack[:-2] + + # Wrap format spec in JoinedStr if it's a constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + else: + # Just pop the value + value = ensure_ast(state.stack[-1]) + new_stack = state.stack[:-1] + format_spec_node = None + + # Determine conversion type from flags + conversion_flags = flags & 0x03 + conversion_map = { + 0: -1, # No conversion + 1: 115, # str (!s) + 2: 114, # repr (!r) + 3: 97, # ascii (!a) + } + conversion = conversion_map[conversion_flags] + + # Create formatted value AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = new_stack + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_313) +@register_handler("FORMAT_SIMPLE", version=PythonVersion.PY_314) +def handle_format_simple( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_SIMPLE formats a string with a single value + # Pops the value and the format string from the stack + assert len(state.stack) >= 1, "Not enough items on stack for FORMAT_SIMPLE" + value = state.stack[-1] + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=None + ) + new_stack = state.stack[:-1] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_313) +@register_handler("FORMAT_WITH_SPEC", version=PythonVersion.PY_314) +def handle_format_with_spec( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # FORMAT_WITH_SPEC formats a value with a format specifier + # Stack order in Python 3.13: format_spec on top, value below + assert len(state.stack) >= 2, "Not enough items on stack for FORMAT_WITH_SPEC" + format_spec = ensure_ast(state.stack[-1]) # Format spec is on top + value = state.stack[-2] # Value is below + + # Check if the value was converted + if isinstance(value, ConvertedValue): + conversion = value.ast_conversion + value = value.value + else: + conversion = -1 + value = ensure_ast(value) + + # Create formatted string AST with specifier + # The format_spec should be wrapped in a JoinedStr if it's a simple constant + if isinstance(format_spec, ast.Constant): + format_spec_node = ast.JoinedStr(values=[format_spec]) + else: + # Already a JoinedStr from nested formatting + assert isinstance(format_spec, ast.JoinedStr) + format_spec_node = format_spec + + formatted_node = ast.FormattedValue( + value=value, conversion=conversion, format_spec=format_spec_node + ) + new_stack = state.stack[:-2] + [ast.JoinedStr(values=[formatted_node])] + return replace(state, stack=new_stack) + + +# ============================================================================ +# CONDITIONAL JUMP HANDLERS +# ============================================================================ + + +def _handle_pop_jump_if( + f_condition: Callable[[ast.expr], ast.expr], + state: ReconstructionState, + instr: dis.Instruction, +) -> ReconstructionState: + # Generic handler for POP_JUMP_IF_* instructions. Pops a value from the + # stack; `condition` is true exactly when the jump is taken. + condition: ast.expr = f_condition(ensure_ast(state.stack[-1])) + + # An inlined-builtin guard is an implementation detail of the interpreter, + # not part of the comprehension: drop it rather than record it as a filter. + if _specialization_guard_edge(state, instr) is not None: + return replace(state, stack=state.stack[:-1]) + + kind, _ = _classify_branch(state, instr) + edge = state.branches.get(instr.offset, BranchEdge.TAKE_JUMP) + + if kind is BranchKind.TERNARY: + return _handle_conditional_expression(state, instr, condition, edge) + + # A filter. The guard is the condition under which *this* path carries on + # toward the element, so it is negated when the path falls through. + guard = condition if edge is BranchEdge.TAKE_JUMP else _negate(condition) + return _attach_filter(state, guard) + + +def _attach_filter( + state: ReconstructionState, guard: ast.expr | None +) -> ReconstructionState: + """Conjoin ``guard`` to the filters of the innermost unfinished comprehension.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if not isinstance(item, CompExp): + continue + + elt: ast.expr = item.value if isinstance(item, ast.DictComp) else item.elt + new_result: CompExp = copy.deepcopy(item) + + if isinstance(elt, Placeholder): + resolved = guard + elif isinstance(elt, ast.IfExp) and any( + isinstance(x, Placeholder) for x in ast.walk(elt) + ): + # A conditional expression was being built up in the element slot, + # but it turned out to be part of this filter's condition. Move it + # back out, plugging the guard into the arm still awaiting a value. + if isinstance(new_result, ast.DictComp): + new_result.key, new_result.value = Placeholder(), Placeholder() + else: + new_result.elt = Placeholder() + + if guard is None: + resolved = None + else: + plugged = ReplacePlaceholder(guard).visit(copy.deepcopy(elt)) + assert isinstance(plugged, ast.expr) + resolved = plugged + else: + continue + + if resolved is not None: + ifs = new_result.generators[-1].ifs + combined = _conjoin(ifs + [resolved]) + assert combined is not None + new_result.generators[-1].ifs = [combined] + + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No comprehension context found for filter condition") + + +def _handle_conditional_expression( + state: ReconstructionState, + instr: dis.Instruction, + condition: ast.expr, + edge: BranchEdge, +) -> ReconstructionState: + """Start an ``ast.IfExp``, marking the arm this path did not take.""" + for pos, item in zip(reversed(range(len(state.stack))), reversed(state.stack)): + if any(isinstance(x, Placeholder) for x in ast.walk(item)): + body: Skipped | Placeholder + orelse: Skipped | Placeholder + skipped = Skipped(id=f".SKIPPED_{instr.offset}") + if edge is BranchEdge.FALL_THROUGH: + body, orelse = skipped, Placeholder() + else: + body, orelse = Placeholder(), skipped + + new_ifexp = ast.IfExp(test=condition, body=body, orelse=orelse) + new_result = ReplacePlaceholder(new_ifexp).visit(copy.deepcopy(item)) + new_stack = state.stack[:pos] + [new_result] + state.stack[pos + 1 : -1] + return replace(state, stack=new_stack) + + raise TypeError("No placeholder found for conditional expression") + + +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_TRUE", version=PythonVersion.PY_314) +def handle_pop_jump_if_true( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_TRUE pops a value from the stack and jumps if it's true + # In Python 3.13, this is used for filter conditions where True means continue + return _handle_pop_jump_if(lambda c: c, state, instr) + + +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_FALSE", version=PythonVersion.PY_314) +def handle_pop_jump_if_false( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_FALSE pops a value from the stack and jumps if it's false + # In comprehensions, this is used for filter conditions + return _handle_pop_jump_if( + lambda c: ast.UnaryOp(op=ast.Not(), operand=c), state, instr + ) + + +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NONE pops a value and jumps if it's None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.Is()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_312) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_313) +@register_handler("POP_JUMP_IF_NOT_NONE", version=PythonVersion.PY_314) +def handle_pop_jump_if_not_none( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # POP_JUMP_IF_NOT_NONE pops a value and jumps if it's not None + return _handle_pop_jump_if( + lambda c: ast.Compare( + left=c, ops=[ast.IsNot()], comparators=[ast.Constant(value=None)] + ), + state, + instr, + ) + + +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_FORWARD", version=PythonVersion.PY_314) +def handle_jump_forward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_FORWARD is used to jump forward in the code + # In generator expressions, this is often used to skip code in conditional logic + return state + + +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD", version=PythonVersion.PY_314) +def handle_jump_backward( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # JUMP_BACKWARD is used to jump back to the beginning of a loop (replaces JUMP_ABSOLUTE in 3.13) + # In generator expressions, this typically indicates the end of the loop body + return state + + +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_BACKWARD_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_backward_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_BACKWARD_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_312) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_313) +@register_handler("JUMP_NO_INTERRUPT", version=PythonVersion.PY_314) +def handle_jump_no_interrupt( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError( + "JUMP_NO_INTERRUPT instruction should not appear in generator comprehensions" + ) + + +@register_handler("JUMP", version=PythonVersion.PY_312) +@register_handler("JUMP", version=PythonVersion.PY_313) +@register_handler("JUMP", version=PythonVersion.PY_314) +def handle_jump( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("JUMP instruction should not appear in generator comprehensions") + + +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_312) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_313) +@register_handler("EXTENDED_ARG", version=PythonVersion.PY_314) +def handle_extended_arg( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # EXTENDED_ARG prefixes an instruction whose argument does not fit in a + # byte. `dis` has already folded it into the following instruction's `arg`, + # so there is nothing left to do here. + return state + + +@register_handler("RESUME", version=PythonVersion.PY_312) +@register_handler("RESUME", version=PythonVersion.PY_313) +@register_handler("RESUME", version=PythonVersion.PY_314) +def handle_resume( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + # RESUME is used for resuming execution after yield/await - mostly no-op for AST reconstruction + return state + + +@register_handler("SEND", version=PythonVersion.PY_312) +@register_handler("SEND", version=PythonVersion.PY_313) +@register_handler("SEND", version=PythonVersion.PY_314) +def handle_send( + state: ReconstructionState, instr: dis.Instruction +) -> ReconstructionState: + raise TypeError("SEND instruction should not appear in generator comprehensions") + + +# ============================================================================ +# UTILITY FUNCTIONS +# ============================================================================ + + +@functools.singledispatch +def ensure_ast(value) -> ast.expr: + """Ensure value is an AST node""" + raise TypeError(f"Cannot convert {type(value)} to AST node") + + +@ensure_ast.register +def _ensure_ast_ast(value: ast.expr) -> ast.expr: + """If already an AST node, return it as is""" + return value + + +@ensure_ast.register(int) +@ensure_ast.register(float) +@ensure_ast.register(str) +@ensure_ast.register(bytes) +@ensure_ast.register(bool) +@ensure_ast.register(complex) +@ensure_ast.register(type(None)) +def _ensure_ast_constant(value) -> ast.Constant: + return ast.Constant(value=value) + + +@ensure_ast.register +def _ensure_ast_tuple(value: tuple) -> ast.Tuple: + """Convert tuple to AST - special handling for dict items""" + if len(value) > 0 and value[0] == "dict_item": + return ast.Tuple( + elts=[ensure_ast(value[1]), ensure_ast(value[2])], ctx=ast.Load() + ) + else: + return ast.Tuple(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +def _unconsumed(value: Iterator) -> typing.Any: + """Return the items an iterator has not yet yielded, as a concrete sequence. + + Built-in sequence iterators pickle as ``(iter, (underlying,), index)``, where + ``index`` is how far the iterator has advanced (absent or ``None`` when it + does not apply). ``reversed`` objects pickle with ``reversed`` as the + callable and count *down* from the end of the underlying sequence. + """ + reduced = value.__reduce__() + assert isinstance(reduced, tuple) and len(reduced) >= 2, ( + f"Cannot recover the contents of {type(value)}" + ) + if not reduced[1]: # an exhausted iterator pickles with no arguments + return () + + underlying = reduced[1][0] + index = reduced[2] if len(reduced) > 2 and reduced[2] is not None else 0 + return underlying[index::-1] if reduced[0] is reversed else underlying[index:] + + +@ensure_ast.register(type(iter((1,)))) +def _ensure_ast_tuple_iterator(value: Iterator) -> ast.Tuple: + return ensure_ast(tuple(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_list(value: list) -> ast.List: + return ast.List(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + + +@ensure_ast.register(type(iter([1]))) +@ensure_ast.register(type(iter({1: 2}.values()))) +@ensure_ast.register(type(iter({1: 2}.items()))) +@ensure_ast.register(type(iter(reversed([1])))) +@ensure_ast.register(reversed) +def _ensure_ast_list_iterator(value: Iterator) -> ast.List: + return ensure_ast(list(_unconsumed(value))) # type: ignore + + +@ensure_ast.register(type(iter("ab"))) # str_ascii_iterator +@ensure_ast.register(type(iter("\xe9b"))) # str_iterator +@ensure_ast.register(type(iter(b"ab"))) +@ensure_ast.register(type(iter(bytearray(b"ab")))) +def _ensure_ast_str_iterator(value: Iterator) -> ast.Constant: + remainder = _unconsumed(value) + # bytearray iteration yields ints, exactly as bytes iteration does + return ensure_ast( # type: ignore + bytes(remainder) if isinstance(remainder, bytearray) else remainder + ) + + +@ensure_ast.register(set) +@ensure_ast.register(frozenset) +def _ensure_ast_set(value: set | frozenset) -> ast.Set: + return ast.Set(elts=[ensure_ast(v) for v in value]) + + +@ensure_ast.register(type(iter({1}))) +def _ensure_ast_set_iterator(value: Iterator) -> ast.Set: + return ensure_ast(set(_unconsumed(value))) # type: ignore + + +@ensure_ast.register +def _ensure_ast_dict(value: dict) -> ast.Dict: + return ast.Dict( + keys=[ensure_ast(k) for k in value.keys()], + values=[ensure_ast(v) for v in value.values()], + ) + + +@ensure_ast.register(type(iter({1: 2}))) +def _ensure_ast_dict_iterator(value: Iterator) -> ast.expr: + return ensure_ast(_unconsumed(value)) + + +@ensure_ast.register(types.BuiltinFunctionType) +@ensure_ast.register(type) +def _ensure_ast_builtin(value: typing.Callable) -> ast.Name: + """A built-in callable is referred to by name, which resolves via builtins. + + Covers both built-in functions (``abs``) and built-in types used as + callables (``bool``), which appear as the predicate of a ``filter`` or the + function of a ``map``. + """ + name = getattr(value, "__name__", None) + assert name and getattr(builtins, name, None) is value, ( + f"Cannot reference non-builtin callable {value!r}" + ) + return ast.Name(id=name, ctx=ast.Load()) + + +@ensure_ast.register(zip) +@ensure_ast.register(enumerate) +@ensure_ast.register(map) +@ensure_ast.register(filter) +def _ensure_ast_iterator_adaptor(value: Iterator) -> ast.Call: + """Rebuild zip/enumerate/map/filter from the arguments they pickle with. + + These wrap other iterators rather than a concrete sequence, so unlike a list + or range iterator they cannot be materialised -- but ``__reduce__`` hands + back their constituent parts, each of which ``ensure_ast`` can handle in + turn. Any already-consumed prefix is reflected in the inner iterators. + """ + reduced = value.__reduce__() + if isinstance(reduced, str): + raise TypeError(f"Cannot convert {type(value)} to AST node") + func, args = reduced[:2] + return ast.Call( + func=ast.Name(id=func.__name__, ctx=ast.Load()), + args=[ensure_ast(arg) for arg in args], + keywords=[], + ) + + +@ensure_ast.register +def _ensure_ast_slice(value: slice) -> ast.Slice: + """A constant slice, as 3.14 emits for `s[1:3]` alongside BINARY_OP/NB_SUBSCR.""" + return ast.Slice( + lower=None if value.start is None else ensure_ast(value.start), + upper=None if value.stop is None else ensure_ast(value.stop), + step=None if value.step is None else ensure_ast(value.step), + ) + + +@ensure_ast.register +def _ensure_ast_range(value: range) -> ast.Call: + return ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ensure_ast(value.start), ensure_ast(value.stop), ensure_ast(value.step)], + keywords=[], + ) + + +@ensure_ast.register(type(iter(range(1)))) +def _ensure_ast_range_iterator(value: Iterator) -> ast.Call: + return ensure_ast(_unconsumed(value)) # type: ignore + + +@ensure_ast.register +def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: + assert inspect.iscode(value), "Input must be a code object" + + name: str = value.co_name.split(".")[-1] + + # Check preconditions + if name in {"", "", "", ""}: + assert name == "" or sys.version_info < (3, 13) + assert name != "" or value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert value.co_argcount == 1 + assert value.co_kwonlyargcount == value.co_posonlyargcount == 0 + assert DummyIterName().id in value.co_varnames + elif name == "": + assert not value.co_flags & inspect.CO_GENERATOR + assert value.co_flags & inspect.CO_NEWLOCALS + assert DummyIterName().id not in value.co_varnames + else: + raise TypeError(f"Unsupported code object type: {value.co_name}") + + # Symbolic execution to reconstruct the AST + result: ast.expr = _symbolic_exec(value) + + # Check postconditions + assert not any(isinstance(x, ast.stmt) for x in ast.walk(result)), ( + "Final return value must not contain statement nodes" + ) + assert not any( + isinstance( + x, + Placeholder + | Skipped + | TargetHole + | CommonConstant + | Null + | CompLambda + | ConvertedValue, + ) + for x in ast.walk(result) + ), "Final return value must not contain temporary nodes" + assert not any(x.arg == ".0" for x in ast.walk(result) if isinstance(x, ast.arg)), ( + "Final return value must not contain .0 argument" + ) + assert not any( + isinstance(x, ast.Name) and x.id == ".0" + for x in ast.walk(result) + if not isinstance(x, DummyIterName) + ), "Final return value must not contain .0 names" + assert sum(1 for x in ast.walk(result) if isinstance(x, DummyIterName)) <= 1, ( + "Final return value must contain at most 1 dummy iterator names" + ) + assert all(x.generators for x in ast.walk(result) if isinstance(x, CompExp)), ( + "Return value must have generators if not a lambda" + ) + + if name == "" and isinstance(result, ast.expr): + # co_varnames lists parameters first: positional, keyword-only, then + # *args and **kwargs if present. Default values are not part of the code + # object -- they are pushed by the caller and attached by MAKE_FUNCTION + # (3.12) or SET_FUNCTION_ATTRIBUTE (3.13). + names = value.co_varnames + n_args, n_kwonly = value.co_argcount, value.co_kwonlyargcount + n_params = n_args + n_kwonly + + vararg = kwarg = None + if value.co_flags & inspect.CO_VARARGS: + vararg = ast.arg(arg=names[n_params]) + n_params += 1 + if value.co_flags & inspect.CO_VARKEYWORDS: + kwarg = ast.arg(arg=names[n_params]) + + args = ast.arguments( + posonlyargs=[ast.arg(arg=arg) for arg in names[: value.co_posonlyargcount]], + args=[ast.arg(arg=arg) for arg in names[value.co_posonlyargcount : n_args]], + vararg=vararg, + kwonlyargs=[ast.arg(arg=arg) for arg in names[n_args : n_args + n_kwonly]], + kw_defaults=[None] * n_kwonly, + kwarg=kwarg, + defaults=[], + ) + return ast.Lambda(args=args, body=result) + elif name == "" and isinstance(result, ast.GeneratorExp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.DictComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.ListComp): + return CompLambda(body=result) + elif name == "" and isinstance(result, ast.SetComp): + return CompLambda(body=result) + else: + raise TypeError(f"Invalid result for type {name}: {result}") + + +@ensure_ast.register +def _ensure_ast_lambda(value: types.LambdaType) -> ast.Lambda: + assert inspect.isfunction(value) and value.__name__.endswith(""), ( + "Input must be a lambda function" + ) + code: types.CodeType = value.__code__ + result = ensure_ast(code) + assert isinstance(result, ast.Lambda), "Lambda body must be an AST Lambda node" + assert not isinstance(result, CompLambda), "Lambda must not be a CompLambda" + return result + + +@ensure_ast.register +def _ensure_ast_genexpr(genexpr: types.GeneratorType) -> ast.GeneratorExp: + assert inspect.isgenerator(genexpr), "Input must be a generator expression" + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must be in created state" + ) + genexpr_ast = ensure_ast(genexpr.gi_code) + assert isinstance(genexpr_ast, CompLambda) + assert genexpr.gi_frame is not None, "Generator must not be exhausted" + geniter_ast = ensure_ast(genexpr.gi_frame.f_locals[".0"]) + result = genexpr_ast.inline(geniter_ast) + assert isinstance(result, ast.GeneratorExp) + assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( + "Generator must stay in created state" + ) + return result + + +# ============================================================================ +# MAIN RECONSTRUCTION FUNCTION +# ============================================================================ + + +def disassemble( + genexpr: Generator[typing.Any, typing.Any, typing.Any], +) -> ast.Expression: + """ + Reconstruct an AST from a generator expression's bytecode. + + This function analyzes the bytecode of a generator object and reconstructs + an abstract syntax tree (AST) that represents the original comprehension + expression. The reconstruction process simulates the Python VM's execution + of the bytecode, building AST nodes instead of executing operations. + + The reconstruction handles complex comprehension features including: + - Multiple nested loops + - Filter conditions (if clauses) + - Complex expressions in the yield/result part + - Tuple unpacking in loop variables + - Various operators and function calls + + Args: + genexpr (Generator[object, None, None]): The generator object to analyze. + Must be a freshly created generator that has not been iterated yet + (in 'GEN_CREATED' state). + + Returns: + ast.Expression: An AST node representing the reconstructed comprehension. + + Raises: + AssertionError: If the input is not a generator or if the generator + has already been started (not in 'GEN_CREATED' state). + + Example: + >>> # Generator expression + >>> g = (x * 2 for x in range(10) if x % 2 == 0) + >>> ast_node = disassemble(g) + >>> isinstance(ast_node, ast.Expression) + True + + >>> # The reconstructed AST can be compiled and evaluated + >>> import ast + >>> code = compile(ast_node, '', 'eval') + >>> result = eval(code) + >>> list(result) + [0, 4, 8, 12, 16] + + Note: + The reconstruction is based on bytecode analysis and may not perfectly + preserve the original source code formatting or variable names in all + cases. However, the semantic behavior of the reconstructed AST should + match the original comprehension. + """ + assert inspect.isgenerator(genexpr), "Input must be a generator expression" + return ast.fix_missing_locations(ast.Expression(ensure_ast(genexpr))) diff --git a/tests/test_internals_disassembler.py b/tests/test_internals_disassembler.py new file mode 100644 index 000000000..55d947f6e --- /dev/null +++ b/tests/test_internals_disassembler.py @@ -0,0 +1,1764 @@ +import ast +import collections.abc +import copy +import typing + +import pytest + +from effectful.internals.disassembly import ( + CompLambda, + DummyIterName, + disassemble, + ensure_ast, +) + + +def compile_and_eval( + node: ast.expr | ast.Expression, globals_dict: dict | None = None +) -> typing.Any: + """Compile an AST node and evaluate it.""" + if globals_dict is None: + globals_dict = {} + + # Wrap in an Expression node if needed + if not isinstance(node, ast.Expression): + node = ast.Expression(body=node) + + # Fix location info + ast.fix_missing_locations(node) + + # Compile and evaluate + code = compile(node, "", "eval") + return eval(code, globals_dict) + + +def materialize[T](genexpr: collections.abc.Generator[T, None, None]) -> list[T]: + """Materialize a nested generator expression to a nested list.""" + + def _materialize(genexpr): + if isinstance(genexpr, str | bytes): + return genexpr + elif isinstance(genexpr, collections.abc.Generator): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, tuple): + # Kept as a tuple so that sets of tuples stay hashable + return tuple(_materialize(item) for item in genexpr) + elif isinstance(genexpr, collections.abc.Sequence): + return [_materialize(item) for item in genexpr] + elif isinstance(genexpr, collections.abc.Set): + return {_materialize(item) for item in genexpr} + elif isinstance(genexpr, collections.abc.Mapping): + return {_materialize(k): _materialize(v) for k, v in genexpr.items()} + else: + return genexpr + + return [_materialize(x) for x in genexpr] + + +def assert_ast_equivalent( + genexpr: collections.abc.Generator[typing.Any, None, None], + reconstructed_ast: ast.AST, + globals_dict: dict | None = None, +): + """Assert that a reconstructed AST produces the same results as the original generator.""" + # Check AST structure + assert isinstance(reconstructed_ast, ast.Expression) + assert hasattr(reconstructed_ast.body, "elt") # The expression part + assert hasattr(reconstructed_ast.body, "generators") # The comprehension part + assert len(reconstructed_ast.body.generators) > 0 + for comp in reconstructed_ast.body.generators: + assert hasattr(comp, "target") # Loop variable + assert hasattr(comp, "iter") # Iterator + assert hasattr(comp, "ifs") # Conditions + + # Save current globals to restore later + curr_globals = globals().copy() + globals().update(globals_dict or {}) + + # Materialize original generator to list for comparison + original_list = materialize(genexpr) + + # Clean up globals to avoid pollution + for key in globals_dict or {}: + if key not in curr_globals: + del globals()[key] + globals().update(curr_globals) + + # Compile and evaluate the reconstructed AST + reconstructed_gen = compile_and_eval(reconstructed_ast, globals_dict) + reconstructed_list = materialize(reconstructed_gen) + assert reconstructed_list == original_list, ( + f"AST produced {reconstructed_list}, expected {original_list}" + ) + + +# ============================================================================ +# BASIC GENERATOR EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple generator expressions + (x for x in range(5)), + (y for y in range(10)), + (item for item in [1, 2, 3]), + # Edge cases for simple generators + (i for i in range(0)), # Empty range + (n for n in range(1)), # Single item range + (val for val in range(100)), # Large range + (x for x in range(-5, 5)), # Negative range + (step for step in range(0, 10, 2)), # Step range + (rev for rev in range(10, 0, -1)), # Reverse range + ], +) +def test_simple_generators(genexpr): + """Test reconstruction of simple generator expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ARITHMETIC AND EXPRESSION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic arithmetic operations + (x * 2 for x in range(5)), + (x + 1 for x in range(5)), + (x - 1 for x in range(5)), + (x**2 for x in range(5)), + (x % 2 for x in range(10)), + (x / 2 for x in range(1, 6)), + (x // 2 for x in range(10)), + # Complex expressions + (x * 2 + 1 for x in range(5)), + ((x + 1) * (x - 1) for x in range(5)), + (x**2 + 2 * x + 1 for x in range(5)), + # Unary operations + (-x for x in range(5)), + (+x for x in range(-5, 5)), + (~x for x in range(5)), + # More complex arithmetic edge cases + (x**3 for x in range(1, 5)), # Higher powers + (x * x * x for x in range(5)), # Repeated multiplication + (x + x + x for x in range(5)), # Repeated addition + (x - x + 1 for x in range(5)), # Operations that might simplify + (x / x for x in range(1, 5)), # Division by self + (x % (x + 1) for x in range(1, 10)), # Modulo with expression + # Nested arithmetic expressions + ((x + 1) ** 2 for x in range(5)), + ((x * 2 + 3) * (x - 1) for x in range(5)), + (x * (x + 1) * (x + 2) for x in range(5)), + # Mixed operations with precedence + (x + 3 * 2 for x in range(3)), + (x * 2 + 9 / 3 for x in range(1, 4)), + ((x + 2) * (x - 2) for x in range(1, 4)), + # Edge cases with zero and one + (x * 0 for x in range(5)), + (x * 1 for x in range(5)), + (x + 0 for x in range(5)), + (x**1 for x in range(5)), + (0 + x for x in range(5)), + (1 * x for x in range(5)), + ], +) +def test_arithmetic_expressions(genexpr): + """Test reconstruction of generators with arithmetic expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FSTRING EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic f-string cases + (f"{x}" for x in range(5)), # Single value, no format + (f"{x} is {x**2}" for x in range(5)), # Multiple values + (f"{x:02d}" for x in range(10)), # Format spec + (f"{x:.2f}" for x in [1.2345, 2.3456, 3.4567]), # Float format spec + # Conversion specifiers + (f"{x!r}" for x in ["hello", "world"]), # repr conversion + (f"{x!s}" for x in [1, 2, 3]), # str conversion + (f"{x!a}" for x in ["hello\n", "world\t"]), # ascii conversion + # Conversion with format spec + (f"{x!r:>10}" for x in ["hello", "world"]), # repr with alignment + (f"{x!s:^15}" for x in [1, 2, 3]), # str with center align + # Empty and literal f-strings + ("" for x in range(3)), # Empty f-string + ("constant" for x in range(3)), # No formatting + (f"x={x}" for x in range(5)), # Literal prefix + (f"result: {x * 2}" for x in range(5)), # Literal with expression + # Complex expressions in f-strings + (f"{x + 1}" for x in range(5)), # Arithmetic + (f"{x * x}" for x in range(5)), # Multiplication + (f"{x % 2}" for x in range(10)), # Modulo + (f"{-x}" for x in range(-2, 3)), # Unary minus + # Nested formatting + (f"{x:0{2}d}" for x in range(5)), # Format spec with expression + (f"{x:>{3 * 2}}" for x in range(5)), # Expression in format spec + # Multiple formatted values + (f"{x} + {y} = {x + y}" for x in range(3) for y in range(3)), # Multiple vars + (f"({x}, {y})" for x in range(2) for y in range(2)), # Tuple display + # F-strings with various data types + (f"{s}" for s in ["hello", "world"]), # Strings + (f"{b}" for b in [True, False]), # Booleans + (f"{n}" for n in [None, None]), # None values + (f"{lst}" for lst in [[1, 2], [3, 4]]), # Lists + # Complex format specifications + (f"{x:+05d}" for x in range(-2, 3)), # Sign, zero pad, width + (f"{x:.2%}" for x in [0.1, 0.25, 0.333]), # Percentage format + (f"{x:.2e}" for x in [100, 1000, 10000]), # Scientific notation + (f"{x:#x}" for x in [10, 15, 255]), # Hex with prefix + (f"{x:b}" for x in [2, 7, 15]), # Binary format + # Edge cases + ("{x}" for x in range(3)), # Escaped braces + (f"{{x}} = {x}" for x in range(3)), # Mixed escaped/formatted + (f"{{{x}}}" for x in range(3)), # Brace around formatted + ], +) +def test_fstring_expressions(genexpr): + """Test reconstruction of generators with f-string expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPARISON OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # All comparison operators + (x for x in range(10) if x < 5), + (x for x in range(10) if x <= 5), + (x for x in range(10) if x > 5), + (x for x in range(10) if x >= 5), + (x for x in range(10) if x == 5), + (x for x in range(10) if x != 5), + # in/not in operators + (x for x in range(10) if x in [2, 4, 6, 8]), + (x for x in range(10) if x not in [2, 4, 6, 8]), + # is/is not operators (with None) + (x for x in [1, None, 3, None, 5] if x is not None), + (x for x in [1, None, 3, None, 5] if x is None), + # Boolean operations - these are complex cases that might need special handling + (x for x in range(10) if not x % 2), + (x for x in range(10) if not (x > 5)), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex comparison edge cases + # Comparisons with expressions + (x for x in range(10) if x * 2 > 10), + (x for x in range(10) if x + 1 <= 5), + (x for x in range(10) if x**2 < 25), + (x for x in range(10) if (x + 1) * 2 != 6), + # Complex membership tests + (x for x in range(20) if x in range(5, 15)), + (x for x in range(10) if x not in range(3, 7)), + (x for x in range(10) if x % 2 in [0]), + (x for x in range(10) if x not in []), # Empty container + # Complex boolean combinations + (x for x in range(20) if not (x < 5 or x > 15)), + (x for x in range(20) if x > 5 and x < 15 and x % 2 == 0), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if not (x > 5 and x < 15)), + # Mixed comparison and boolean operations + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if not (x % 2 == 0 and x % 3 == 0)), + # Edge cases with identity comparisons + (x for x in [0, 1, 2, None, 4] if x is not None and x > 1), + (x for x in [True, False, 1, 0] if x is True), + (x for x in [True, False, 1, 0] if x is not False), + ], +) +def test_comparison_operators(genexpr): + """Test reconstruction of all comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CHAINED COMPARISON TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Chained comparisons + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + ], +) +def test_chained_comparison_operators(genexpr): + """Test reconstruction of chained (ternary) comparison operators.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# FILTERED GENERATOR TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple filters + (x for x in range(10) if x % 2 == 0), + (x for x in range(10) if x > 5), + (x for x in range(10) if x < 5), + (x for x in range(10) if x != 5), + # Complex filters + (x for x in range(20) if x % 2 == 0 if x % 3 == 0), + (x for x in range(100) if x > 10 if x < 90 if x % 5 == 0), + # Filters with expressions + (x * 2 for x in range(10) if x % 2 == 0), + (x**2 for x in range(10) if x > 3), + # Boolean operations in filters + (x for x in range(10) if not x % 2), + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(10) if x < 3 or x > 7), + # More complex filter edge cases + (x for x in range(50) if x % 7 == 0), # Different modulo + (x for x in range(10) if x >= 0), # Always true condition + (x for x in range(10) if x < 0), # Always false condition + ( + x for x in range(20) if x % 2 == 0 and x % 3 == 0 + ), # Multiple conditions with and + ( + x for x in range(20) if x % 2 == 0 or x % 3 == 0 + ), # Multiple conditions with or + # Nested boolean operations + (x for x in range(20) if (x > 5 and x < 15) or x == 0), + (x for x in range(20) if not (x > 10 and x < 15)), + (x for x in range(50) if x > 10 and (x % 2 == 0 or x % 3 == 0)), + # Multiple consecutive filters + (x for x in range(100) if x > 20 if x < 80 if x % 10 == 0), + (x for x in range(50) if x % 2 == 0 if x % 3 != 0 if x > 10), + # Filters with complex expressions + (x + 1 for x in range(20) if (x * 2) % 3 == 0), + (x**2 for x in range(10) if x * (x + 1) > 10), + (x / 2 for x in range(1, 20) if x % (x // 2 + 1) == 0), + # Edge cases with truthiness + (x for x in range(10) if x), # Truthy filter + (x for x in range(-5, 5) if not x), # Falsy filter + (x for x in range(10) if bool(x % 2)), # Explicit bool conversion + ], +) +def test_filtered_generators(genexpr): + """Test reconstruction of generators with if conditions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# NESTED LOOP TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic nested loops + ((x, y) for x in range(3) for y in range(3)), + (x + y for x in range(3) for y in range(3)), + (x * y for x in range(1, 4) for y in range(1, 4)), + # Nested with filters + ((x, y) for x in range(5) for y in range(5) if x < y), + (x + y for x in range(5) if x % 2 == 0 for y in range(5) if y % 2 == 1), + # Triple nested + (x + y + z for x in range(2) for y in range(3) for z in range(4)), + ((x, y, z) for x in range(2) for y in range(3) for z in range(4)), + # More complex nested loop edge cases + # Different sized ranges + ((x, y) for x in range(2) for y in range(5)), + ((x, y) for x in range(10) for y in range(2)), + # Asymmetric operations + (x - y for x in range(5) for y in range(3)), + (x / (y + 1) for x in range(1, 6) for y in range(3)), + (x**y for x in range(1, 4) for y in range(3)), + # Complex expressions with multiple variables + (x * y + x for x in range(3) for y in range(3)), + (x + y + x * y for x in range(1, 4) for y in range(1, 4)), + ((x + y) ** 2 for x in range(3) for y in range(3)), + # Filters on different loop levels + ((x, y) for x in range(10) if x % 2 == 0 for y in range(10) if y % 3 == 0), + (x * y for x in range(5) for y in range(5) if x != y), + (x + y for x in range(5) for y in range(5) if x + y < 5), + # Triple and quadruple nested with various patterns + (x + y + z for x in range(2) for y in range(2) for z in range(2)), + (x * y * z for x in range(1, 3) for y in range(1, 3) for z in range(1, 3)), + ( + (x, y, z, w) + for x in range(2) + for y in range(2) + for z in range(2) + for w in range(2) + ), + # Nested loops with complex filters + ((x, y) for x in range(5) if x > 1 for y in range(5) if x < y), + (x + y for x in range(3) if x > 0 for y in range(3)), + # Mixed range types + ((x, y) for x in range(-2, 2) for y in range(0, 4, 2)), + (x * y for x in range(5, 0, -1) for y in range(1, 6)), + # Dependent nested loops + ((x, y) for x in range(3) for y in range(x, 3)), + (x + y for x in range(3) for y in range(x + 1, 3)), + ], +) +def test_nested_loops(genexpr): + """Test reconstruction of generators with nested loops.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# =========================================================================== +# NESTED COMPREHENSIONS +# =========================================================================== + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested generators + ((x for x in range(i + 1)) for i in range(5)), + ((x for j in range(i) for x in range(j)) for i in range(5)), + (((x for x in range(i + j)) for j in range(i)) for i in range(5)), + # nested generators with filters + ((x for x in range(i)) for i in range(5) if i > 0), + ((x for x in range(i) if x < i) for i in range(5) if i > 0), + (((x for x in range(i + j) if x < i + j) for j in range(i)) for i in range(5)), + # aggregation function call + (sum(x for x in range(i + 1)) for i in range(3)), + (max(x for x in range(i + 1)) for i in range(3)), + (dict((x, x + 1) for x in range(i + 1)) for i in range(3)), + (set(x for x in range(i + 1)) for i in range(3)), + # map + (list(map(abs, (x + 1 for x in range(i + 1)))) for i in range(3)), + (list(enumerate(x + 1 for x in range(i + 1))) for i in range(3)), + # nesting on both sides + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ((y for y in range(x)) for x in (x_ + 1 for x_ in range(5))), + ], +) +def test_nested_comprehensions(genexpr): + """Test reconstruction of nested comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_nested_comprehensions_multiline(): + """The same filter reconstructs the same way however the source is laid out. + + On Python 3.12 these two spellings disassemble to different jump layouts -- + only the one-line form emits POP_JUMP_IF_TRUE -- which used to make the + multiline form come out negated. + """ + one_line = (x for x in range(5) if x > 1) + assert_ast_equivalent(one_line, disassemble(one_line)) + + multiline = ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + assert_ast_equivalent(multiline, disassemble(multiline)) + + assert ast.unparse(disassemble(x for x in range(5) if x > 1)) == ast.unparse( + disassemble( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ) + ) + + +# ============================================================================ +# DIFFERENT COMPREHENSION TYPES +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Comprehensions as iterator constants + (x_ for x_ in [x for x in range(5)]), + (x_ for x_ in {x for x in range(5)}), + (x_ for x_ in {x: x**2 for x in range(5)}), + # Comprehensions as yield expressions + ([y * 2 for y in range(x + 1)] for x in range(3)), + ({y + 3 for y in range(x + 1)} for x in range(3)), + ({y: y**2 for y in range(x + 1)} for x in range(3)), + # nested non-generators + ([x for x in range(i)] for i in range(5)), + ([x for j in range(i) for x in range(j)] for i in range(5)), + ({x: x**2 for x in range(i)} for i in range(5)), + ([[x for x in range(i + j)] for j in range(i)] for i in range(5)), + # Nested comprehensions with filters inside + ([x for x in range(i)] for i in range(5) if i > 0), + ([x for x in range(i) if x < i] for i in range(5) if i > 0), + ([[x for x in range(i + j) if x < i + j] for j in range(i)] for i in range(5)), + ], +) +def test_different_comprehension_types(genexpr): + """Test reconstruction of different comprehension types.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# CONDITIONAL EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # simple conditional expressions without nesting + ((lambda x: x if x % 2 == 0 else -x)(xi) for xi in range(5)), + ((lambda x: (x + 1) if x < 5 else (x - 1))(xi) for xi in range(10)), + ((lambda x: (x * 2) if x > 0 else (x / 2))(xi) for xi in range(-5, 5)), + ((lambda x: (x**2) if x != 0 else 1)(xi) for xi in range(-3, 4)), + # simple conditional expressions with negation + ((lambda x: (x + 10) if not (x < 5) else (x - 10))(xi) for xi in range(20)), + ((lambda x: (x * 3) if not (x % 2 == 0) else (x // 3))(xi) for xi in range(10)), + ((lambda x: (x**3) if not (x < 0) else (x**0.5))(xi) for xi in range(-5, 15)), + # conditional expressions with lazy test + ( + (lambda x: (x + 10) if (x > 5 and x < 15) else (x - 10))(xi) + for xi in range(20) + ), + ( + (lambda x: (x * 3) if (x % 2 == 0 or x % 3 == 0) else (x // 3))(xi) + for xi in range(10) + ), + ( + (lambda x: (x**3) if not (x < 0 or x > 10) else (x**0.5))(xi) + for xi in range(-5, 15) + ), + ], +) +def test_conditional_expressions_simple_no_comprehension(genexpr): + """Test reconstruction of simple conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # nested conditional expressions + ( + (lambda x: (x + 1) if x < 5 else ((x - 1) if x < 10 else (x * 2)))(xi) + for xi in range(15) + ), + ( + ( + lambda x: ( + (x * 2) if x % 2 == 0 else ((x // 2) if x % 3 == 0 else (x + 2)) + ) + )(xi) + for xi in range(10) + ), + ( + (lambda x: (x**2) if x > 0 else ((-x) ** 2 if x < -5 else 1))(xi) + for xi in range(-10, 5) + ), + ], +) +def test_conditional_expressions_nested_no_comprehension(genexpr): + """Test reconstruction of nested conditional expressions isolated from comprehension bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Basic conditional expressions in comprehension bodies + ((x if x % 2 == 0 else -x) for x in range(5)), + ((x * 2 if x > 0 else x / 2) for x in range(-3, 4)), + ((x**2 if x != 0 else 1) for x in range(-2, 3)), + # Conditional expressions with filters + ((x if x % 2 == 0 else -x) for x in range(10) if x > 2), + ((x * 3 if x > 5 else x + 1) for x in range(20) if x % 3 == 0), + # Nested loops with conditional expressions + ((x + y if x > y else x - y) for x in range(3) for y in range(3)), + ( + (x * y if x != 0 and y != 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Multiple conditional expressions + ( + (x if x > 0 else 0) + (y if y > 0 else 0) + for x in range(-2, 3) + for y in range(-2, 3) + ), + # Conditional expressions in different parts + ([x if x > 0 else -x for x in range(i)] for i in range(1, 4)), + ((x if x % 2 == 0 else -x) for x in (y if y > 2 else y + 10 for y in range(5))), + # Complex nested conditional expressions + ((x if x > 0 else (x + 5 if x > -3 else x * 2)) for x in range(-5, 5)), + ((x * 2 if x > 0 else (x / 2 if x < 0 else 1)) for x in range(-3, 4)), + # Conditional expressions with function calls + ((abs(x) if x < 0 else x) for x in range(-3, 4)), + ((max(x, 0) if x is not None else 0) for x in [None, -1, 0, 1, 2]), + # Mixed with other complex expressions + ((x + 1 if x % 2 == 0 else x - 1) * 2 for x in range(5)), + ((x, y, x + y if x > y else x - y) for x in range(3) for y in range(3)), + ], +) +def test_conditional_expressions_simple_comprehensions(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # `and` chains compile to a run of jumps that each fall through to the + # loop back-edge, which the disassembler folds into a single filter. + (x for x in range(10) if x > 2 and x < 8), + (x for x in range(20) if x > 5 and x % 2 == 0 and x < 15), + (x for x in range(-10, 10) if abs(x) > 3 and x % 2 == 0), + (x for x in ["hello", "world", "test"] if len(x) > 3 and x.startswith("h")), + (x for x in range(20) if x % 2 == 0 and x % 3 == 0 and x > 0 and x < 18), + ((x, y) for x in range(5) for y in range(5) if x < y and x + y > 2), + # `or` in filter position: reconstructed incorrectly, see the marker. + (x for x in range(10) if x < 3 or x > 7), + (x for x in range(20) if x < 5 or x > 15 or x == 10), + (x for x in range(20) if (x > 10 and x % 2 == 0) or (x < 5 and x % 3 == 0)), + (x for x in range(20) if x > 5 and (x < 10 or x > 15)), + (x for x in range(100) if (x > 10 and x < 50) and (x % 3 == 0 or x % 5 == 0)), + # `not (a and b)` is compiled exactly like `not a or not b`. + (x for x in range(100) if not (x > 30 and x < 70)), + # Chained comparisons in filter position. + (x for x in range(20) if 5 < x < 15), + (x for x in range(20) if 0 <= x <= 10), + (x for x in range(50) if 10 < x < 20 < x * 2), + (x for x in range(10) if 0 <= x <= 5 <= x + 5), + (x for x in range(50) if 5 < x < 15 and x % 2 == 0), + (x for x in range(50) if x > 20 or 5 < x < 15), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_filters(genexpr): + """Lazy boolean operators and chained comparisons in *filter* position. + + This is the hard case: a filter's condition is recognised structurally, by + the jump falling through to the loop back-edge, so any condition CPython + compiles with an intermediate join point is misread. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # The same operators in *ternary* position all work: both arms produce a + # value, so the fork/merge machinery in _symbolic_exec applies directly. + ((x if x > 5 and x < 15 else 0) for x in range(20)), + ((x if x < 3 or x > 17 else -x) for x in range(20)), + ((x if 5 < x < 15 else 0) for x in range(20)), + ((x * 2 if 0 <= x <= 10 else x / 2) for x in range(-5, 15)), + ((x if x > 2 and x < 8 else -x) for x in range(10)), + ((x if x < 2 or x > 8 else -x) for x in range(10)), + ((x if not (x > 2 and x < 8) else -x) for x in range(10)), + ((x if 0 <= x <= 5 <= x + 5 else -x) for x in range(10)), + ((x if x > 1 and x < 9 or x == 0 else -x) for x in range(10)), + # ... including nested inside another ternary + ((x if x > 5 or x < 2 else (0 if x == 3 else 1)) for x in range(10)), + ((x if x > 5 else (0 if 2 < x < 4 else 1)) for x in range(10)), + ], +) +def test_lazy_boolean_and_chained_comparisons_in_ternaries(genexpr): + """Lazy boolean operators and chained comparisons in conditional expressions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_short_circuit_filter_is_a_disjunction_of_paths(): + """A short-circuiting filter is rebuilt as the disjunction over its paths. + + `x < 3 or x > 7` used to be misread as a conditional expression, yielding an + AST that compiled and ran but computed `[0, 1, 2, False, False, ...]` instead + of `[0, 1, 2, 8, 9]`. Each path through the condition now contributes one + disjunct, so the reconstruction is equivalent rather than merely plausible. + """ + genexpr = (x for x in range(10) if x < 3 or x > 7) + reconstructed = disassemble(genexpr) + + assert isinstance(reconstructed.body, ast.GeneratorExp) + filters = reconstructed.body.generators[0].ifs + assert len(filters) == 1 + assert isinstance(filters[0], ast.BoolOp) and isinstance(filters[0].op, ast.Or) + assert materialize(compile_and_eval(reconstructed)) == [0, 1, 2, 8, 9] + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple conditional as function argument + (max(x if x > 0 else 0, 1) for x in range(-2, 3)), + (abs(x if x < 0 else -x) for x in range(-3, 3)), + (len(str(x) if x > 10 else "small") for x in range(15)), + # Multiple conditional arguments + ( + max(x if x > 0 else 0, y if y > 0 else 0) + for x in range(-1, 2) + for y in range(-1, 2) + ), + ( + pow(x if x != 0 else 1, y if y > 0 else 1) + for x in range(3) + for y in range(3) + ), + # Nested function calls with conditionals + (max(abs(x if x < 0 else -x), 1) for x in range(-3, 4)), + (int(str(x if x > 5 else x + 10)) for x in range(10)), + # Conditionals in keyword arguments (using dict constructor as example) + (dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), + # Method calls with conditional arguments + ([1, 2, 3].index(x if x in [1, 2, 3] else 1) for x in range(5)), + ("hello".replace("l", x if isinstance(x, str) else "X") for x in ["a", 1, "b"]), + # Complex nested case: conditional in function argument, function call in conditional + (abs(x if len(str(x)) > 1 else x * 10) for x in range(15)), + # Mixed: conditional in function call within comprehension filter + (x for x in range(20) if max(x if x > 10 else 0, 5) > 8), + ], +) +def test_conditional_expressions_function_arguments(genexpr): + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# GENERATOR EXPRESSION WITH GLOBALS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using constants + ((x + a for x in range(5)), {"a": 10}), # type: ignore # noqa: F821 + ((data[i] for i in range(2)), {"data": [3, 4]}), # type: ignore # noqa: F821 + # Using global functions + ((abs(x) for x in range(-5, 5)), {"abs": abs}), + ((len(s) for s in ["a", "ab", "abc"]), {"len": len}), + ((max(x, 5) for x in range(10)), {"max": max}), + ((min(x, 5) for x in range(10)), {"min": min}), + ((round(x / 3, 2) for x in range(10)), {"round": round}), + ], +) +def test_variable_lookup(genexpr, globals_dict): + """Test reconstruction of expressions with globals.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# EDGE CASES AND COMPLEX SCENARIOS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + # Using lambdas and functions + (((lambda y: y * 2)(x) for x in range(5)), {}), + (((lambda y: y + 1)(x) for x in range(5)), {}), + (((lambda y: y**2)(x) for x in range(5)), {}), + (((lambda a, b: a + b)(x, x) for x in range(5)), {}), + (((lambda: (x for x in range(i)))() for i in range(3)), {}), + ((f(x) for x in range(5)), {"f": lambda y: y * 3}), # type: ignore # noqa: F821 + # Attribute access + ((x.real for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.imag for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + ((x.conjugate() for x in [1 + 2j, 3 + 4j, 5 + 6j]), {}), + # slicing and indexing + ((s[:2] for s in ["hello", "world"]), {}), + ((s[1:3] for s in ["hello", "world"]), {}), + ((s[-1] for s in ["hello", "world"]), {}), + ((s[0:3] for s in ["hello", "world"]), {}), + ((s[::-1] for s in ["hello", "world"]), {}), + ((s[1:2:] for s in ["hello", "world"]), {}), + # Method calls + ((s.upper() for s in ["hello", "world"]), {}), + ((s.lower() for s in ["HELLO", "WORLD"]), {}), + ((s.strip() for s in [" hello ", " world "]), {}), + ((x.bit_length() for x in range(1, 10)), {}), + ((str(x).zfill(3) for x in range(10)), {"str": str}), + # Subscript operations + (((10, 20, 30)[i] for i in range(3)), {}), + (([10, 20, 30][i] for i in range(3)), {}), + (({"a": 1, "b": 2, "c": 3}[k] for k in ["a", "b", "c"]), {}), + (("hello"[i] for i in range(5)), {}), + ((data[i][j] for i in range(2) for j in range(2)), {"data": [[1, 2], [3, 4]]}), # type: ignore # noqa: F821 + # # More complex attribute chains + # ((obj.value.bit_length() for obj in [type('', (), {'value': x})() for x in range(1, 5)]), {}), + # Multiple function calls + ((abs(max(x, -x)) for x in range(-3, 4)), {"abs": abs, "max": max}), + ((len(str(x)) for x in range(100, 110)), {"len": len, "str": str}), + # Mixed operations + ( + (abs(x) + len(str(x)) for x in range(-10, 10)), + {"abs": abs, "len": len, "str": str}, + ), + ((s.upper().lower() for s in ["Hello", "World"]), {}), + # Edge cases with complex data structures + (((1, 2, 3)[x % 3] for x in range(10)), {}), + (([1, 2, 3][x % 3] for x in range(10)), {}), + (({1, 2, 3} for x in range(10)), {}), + # (({"even": x, "odd": x + 1}["even" if x % 2 == 0 else "odd"] for x in range(5)), {}), + # Function calls with multiple arguments + ((pow(x, 2, 10) for x in range(5)), {"pow": pow}), + ((divmod(x, 3) for x in range(10)), {"divmod": divmod}), + ], +) +def test_complex_scenarios(genexpr, globals_dict): + """Test reconstruction of complex generator expressions.""" + ast_node = disassemble(genexpr) + + # Need to provide the same globals for evaluation + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +# ============================================================================ +# UNPACKING LOOP TARGETS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Simple tuple targets + ((a, b) for a, b in [(1, 2), (3, 4)]), + (a + b for a, b in [(1, 2), (3, 4)]), + (a * b for a, b in [(2, 3), (4, 5)]), + ((b, a) for a, b in [(1, 2), (3, 4)]), + ((a, b, c) for a, b, c in [(1, 2, 3), (4, 5, 6)]), + ((a, b, c, d) for a, b, c, d in [(1, 2, 3, 4)]), + # Nested tuple targets + ((a, b, c) for a, (b, c) in [(1, (2, 3)), (4, (5, 6))]), + ((a, b, c) for (a, b), c in [((1, 2), 3)]), + ((a, b, c, d) for (a, b), (c, d) in [((1, 2), (3, 4))]), + ((a, b, c) for a, (b, (c,)) in [(1, (2, (3,)))]), + # Unpacking over dict views + ((k, v) for k, v in {1: "a", 2: "b"}.items()), + (v for k, v in {1: "a", 2: "b"}.items()), + # Unpacking combined with filters + ((a, b) for a, b in [(1, 2), (3, 1)] if a < b), + (a + b for a, b in [(1, 2), (3, 4)] if a % 2 == 0), + ((a, b) for a, b in [(1, 2), (3, 4)] if a > 0 if b > 3), + # Unpacking in nested loops, in either position + ((x, a, b) for x in range(2) for a, b in [(1, 2), (3, 4)]), + ((a, b, y) for a, b in [(1, 2)] for y in range(2)), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4), (5, 6)]), + ((a, b, x) for a, b in [(1, 2), (3, 4)] for x in range(a)), + # Unpacking inside other comprehension types + ([a for a, b in [(1, 2), (3, 4)]] for _ in range(2)), + ({a for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ({a: b for a, b in [(1, 2), (3, 4)]} for _ in range(2)), + ((a for a, b in [(1, 2), (3, 4)]) for _ in range(2)), + # Unpacking with a conditional expression in the body + ((a if a > b else b) for a, b in [(1, 2), (4, 3)]), + ], +) +def test_unpacking_targets(genexpr): + """Test reconstruction of comprehensions that unpack their loop target.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + ((x, a, b) for x in range(2) for a, b in [(1, 2)]), + ((a, b, c, d) for a, b in [(1, 2)] for c, d in [(3, 4)]), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3)]), + ((x, a, b, c) for x in range(2) for a, (b, c) in [(1, (2, 3))]), + ((x, a) for x in range(2) for (a,) in [(1,)]), + ((x, a, b) for x in range(2) for a, *b in [(1,)]), # type: ignore[var-annotated] + ((x, y, a) for x in range(2) for y in range(2) for a, b in [(1, 2)]), + ], +) +def test_unpacking_over_single_element_literal(genexpr): + """A one-element inner loop over a literal. + + Python 3.14 unrolls this: it assigns the targets outright and emits no + FOR_ITER, so the loop is not there to be recovered. The names it bound are + substituted at their uses instead, which reproduces the same elements. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred last: UNPACK_EX with only a "before" count + (a for a, *b in [(1, 2, 3), (4, 5, 6)]), + (b for a, *b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for a, *b in [(1, 2, 3)]), + ((a, b, c) for a, b, *c in [(1, 2, 3, 4)]), + # Starred first: the "after" count lives in the high byte of the + # argument, so the instruction is prefixed with EXTENDED_ARG + (a for *a, b in [(1, 2, 3), (4, 5, 6)]), + (b for *a, b in [(1, 2, 3), (4, 5, 6)]), + ((a, b) for *a, b in [(1, 2, 3)]), + # Starred in the middle + ((a, b, c) for a, *b, c in [(1, 2, 3, 4)]), + ((a, b, c) for a, *b, c in [(1, 2, 3, 4, 5)]), + # Starred target that collects nothing + ((a, b) for a, *b in [(1,)]), # type: ignore[var-annotated] + # Combined with filters, nesting and other comprehension types + ((a, b) for a, *b in [(1, 2), (3, 4)] if a > 1), + ((x, a, b) for x in range(2) for a, *b in [(1, 2, 3), (4, 5, 6)]), + ([a for a, *b in [(1, 2, 3)]] for _ in range(2)), + ], +) +def test_unpacking_starred_targets(genexpr): + """Test reconstruction of starred loop targets (UNPACK_EX).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# OUTERMOST ITERABLE TYPES +# +# The outermost iterable is not part of the comprehension's bytecode: it is a +# live object reachable through `gi_frame.f_locals[".0"]`, so `ensure_ast` has +# to rebuild an expression for it from the object alone. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Strings and bytes + (c for c in "hello"), + (c.upper() for c in "hello" if c != "l"), + (c for c in "h\xe9llo"), # non-ASCII takes a different iterator type + (b for b in b"abc"), + (b for b in bytearray(b"abc")), + # Sequences + (x for x in [1, 2, 3]), + (x for x in (1, 2, 3)), + (x for x in range(3)), + # Sets and frozensets + (x for x in {1, 2, 3}), + (x for x in frozenset({1, 2, 3})), + # Dict views + (k for k in {1: "a", 2: "b"}), + (k for k in {1: "a", 2: "b"}.keys()), + (v for v in {1: "a", 2: "b"}.values()), + (kv for kv in {1: "a", 2: "b"}.items()), + # reversed() over each of the underlying sequence types + (x for x in reversed([1, 2, 3])), + (x for x in reversed((1, 2, 3))), + (c for c in reversed("abc")), + (x for x in reversed(range(3))), + # Comprehensions as the outermost iterable + (x for x in (y for y in range(3))), + (x for x in [y for y in range(3)]), + (x for x in {y for y in range(3)}), + # Nested/structured contents + (t for t in [(1, 2), (3, 4)]), + (d for d in [{"a": 1}, {"b": 2}]), + (x for x in [[1, 2], [3, 4]]), + ], +) +def test_outermost_iterable_types(genexpr): + """Test reconstruction of the outermost iterable from the live object.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # zip/enumerate/map/filter wrap other iterators rather than a concrete + # sequence, but pickle with their constituent parts. + (x for x in zip([1, 2], [3, 4])), + ((a, b) for a, b in zip([1, 2], [3, 4])), + (x for x in zip("ab", range(2), [7, 8])), + (x for x in enumerate("ab")), + ((i, c) for i, c in enumerate("abc")), + (x for x in enumerate("ab", 5)), + (x for x in map(abs, [-1, 2])), + (x for x in map(max, [1, 2], [3, 0])), + (x for x in filter(None, [0, 1, 2])), + (x for x in filter(bool, [0, 1, 2])), + # Nested adaptors, and adaptors over non-sequence iterables + (x for x in zip(range(2), map(abs, [-1, -2]))), + (x for x in enumerate(filter(None, [0, 1]))), + (x for x in map(abs, range(-2, 2))), + (x for x in zip("ab", (y for y in range(2)))), + # With a filter and a non-trivial element expression + (a + b for a, b in zip([1, 2], [3, 4]) if a > 1), + ], +) +def test_outermost_iterable_adaptors(genexpr): + """zip/enumerate/map/filter are rebuilt from the parts they pickle with.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A conditional expression as an inner loop's iterable. Its value is + # consumed by FOR_ITER rather than by the yield, so the half-built + # IfExp has to be pulled back out of the element slot. + (y for x in range(4) for y in (range(x) if x % 2 == 0 else range(x, x + 2))), + (y for x in range(4) for y in ([x] if x else [0])), + (y for x in range(4) for y in (range(x) if x > 1 else range(1))), + # An *empty* list literal as an arm is misread: BUILD_LIST(0) is also + # how an inlined list comprehension starts, and nothing later + # disambiguates the two here. + pytest.param( + (y for x in range(4) for y in ([x] if x else [])), + marks=pytest.mark.xfail( + strict=True, + reason="an empty list literal is indistinguishable from the start of a list comprehension", + ), + ), + ((x, y) for x in range(3) for y in ([0] if x % 2 else [1, 2])), + # ... with a filter on the inner loop, and nested two deep + (y for x in range(4) for y in (range(x) if x % 2 == 0 else [9]) if y > 0), + pytest.param( + ( + z + for x in range(3) + for y in (range(x) if x else [0]) + for z in ([y] if y else [7]) + ), + marks=pytest.mark.xfail( + strict=True, + reason="two conditional iterables in one comprehension leave paths that do not pairwise merge", + ), + ), + # ... and one whose arms are comprehensions of different kinds + (y for x in range(3) for y in ([i for i in range(x)] if x else {8})), + ], +) +def test_conditional_expression_as_iterable(genexpr): + """Test a conditional expression in the iterable position of a for clause.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # An always-false filter lets the compiler drop the body entirely, so no + # element expression survives in the bytecode. The loop still runs. + (x for x in range(6) if False), + (x for x in range(6) if x and False), + (y for x in range(6) if False and (y := x)), # noqa: F821 + ([x] for x in range(4) if False), + ((x, y) for x in range(4) for y in range(3) if False), + (x for x in range(6) if False if x > 1), + ({x for x in range(3)} for _ in range(2) if False), + ], +) +def test_unreachable_comprehension_body(genexpr): + """A comprehension whose body the compiler proved unreachable yields nothing.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + assert materialize(compile_and_eval(ast_node)) == [] + + +def test_outermost_iterable_partially_consumed_adaptor(): + """A consumed prefix is reflected in the adaptor's inner iterators.""" + zipped = zip([1, 2, 3], [4, 5, 6]) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert materialize(genexpr) == [7, 9] + assert materialize(compile_and_eval(reconstructed)) == [7, 9] + + +def test_outermost_iterable_partially_consumed(): + """Only the *unconsumed* remainder of the outermost iterator belongs in the AST.""" + iterator = iter([10, 20, 30, 40]) + next(iterator) + next(iterator) + + genexpr = (x + 1 for x in iterator) + assert ast.unparse(disassemble(genexpr)) == "(x + 1 for x in [30, 40])" + assert materialize(genexpr) == [31, 41] + + +def test_outermost_iterable_partially_consumed_str(): + iterator = iter("hello") + next(iterator) + + genexpr = (c for c in iterator) + assert ast.unparse(disassemble(genexpr)) == "(c for c in 'ello')" + assert materialize(genexpr) == ["e", "l", "l", "o"] + + +# ============================================================================ +# BINARY OPERATORS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # Bitwise operators, which BINARY_OP folds in with the arithmetic ones + (x & 3 for x in range(8)), + (x | 3 for x in range(8)), + (x ^ 3 for x in range(8)), + (x << 2 for x in range(4)), + (x >> 1 for x in range(8)), + (~x & 7 for x in range(8)), + # Mixed precedence across the whole operator table + (x & 1 | x >> 2 ^ 3 for x in range(8)), + ((x | 1) & (x ^ 2) for x in range(8)), + (x + 1 & x - 1 for x in range(8)), + (x * 2 % 5 // 2 for x in range(8)), + (x**2 - x // 2 + x % 3 for x in range(1, 8)), + # Operators on non-numeric operands + (s + "!" for s in ["a", "b"]), + (s * 2 for s in ["a", "b"]), + (t + (9,) for t in [(1,), (2,)]), + (frozenset({x}) | frozenset({9}) for x in range(3)), + (frozenset({x, 1}) & frozenset({1}) for x in range(3)), + (frozenset({x, 1}) ^ frozenset({1}) for x in range(3)), + ({"a": x} | {"b": 0} for x in range(3)), + ], +) +def test_binary_operators(genexpr): + """Test reconstruction of the full BINARY_OP table.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_matmul_operator(): + """BINARY_OP argument 4 is `@`, which no built-in type implements. + + The comprehension is disassembled but never evaluated, so this checks the + reconstructed source rather than the reconstructed values. + """ + genexpr = (a @ b for a in [1, 2]) # noqa: F821 + assert ast.unparse(disassemble(genexpr)) == "(a @ b for a in (1, 2))" + + +# ============================================================================ +# KEYWORD ARGUMENTS AT CALL SITES +# +# Python 3.12 compiles these as KW_NAMES followed by CALL; 3.13 replaced the +# pair with a single CALL_KW instruction. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr,globals_dict", + [ + ((dict(a=x) for x in range(3)), {}), + ((dict(a=x, b=x * 2) for x in range(3)), {}), + ((dict(a=x if x > 0 else 0, b=x * 2 if x < 5 else x) for x in range(8)), {}), + # Mixed positional and keyword arguments + ((sorted([3, x], reverse=True) for x in range(3)), {}), + ((sorted([x, 1], key=abs) for x in range(3)), {}), + ((sorted([x, 1], key=abs, reverse=True) for x in range(3)), {}), + ((int(str(x), base=8) for x in range(8)), {}), + ((round(x / 3, ndigits=2) for x in range(5)), {}), + # Keyword arguments on a method call + (("a,b".split(sep=",") for x in range(2)), {}), + (("a-b".replace("-", "+") for x in range(2)), {}), + # Nested calls, each with keywords + ((dict(a=dict(b=x)) for x in range(3)), {}), + ((sorted(sorted([x, 1]), reverse=True) for x in range(3)), {}), + # Keyword arguments to a user-supplied callable + ((f(x, scale=2) for x in range(3)), {"f": lambda v, scale=1: v * scale}), # type: ignore # noqa: F821 + ], +) +def test_keyword_arguments(genexpr, globals_dict): + """Test reconstruction of calls with keyword arguments.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node, globals_dict) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Starred positional arguments + (max(*[x, 1]) for x in range(3)), + (max(1, *[x, 2]) for x in range(3)), + (max(*[x, 1], *[2, 3]) for x in range(3)), + (max(*(x, 1)) for x in range(3)), + (sum([x, 1], *[0]) for x in range(3)), + # Double-starred keyword arguments + (dict(**{"a": x}) for x in range(3)), + (dict(a=x, **{"b": 1}) for x in range(3)), + (dict(**{"a": x}, **{"b": 2}) for x in range(3)), + (sorted([x, 1], **{"reverse": True}) for x in range(3)), # type: ignore[call-overload] + # Both at once + (max(*[[x, 1]], **{"default": 0}) for x in range(3)), # type: ignore[call-overload] + # On a method call, where the callable comes with a `self` + ("-".join(*[[str(x), "z"]]) for x in range(3)), + # Unpacking a comprehension, and unpacking into a nested call + (max(*[y for y in range(x + 2)]) for x in range(3)), + (max(*[abs(y) for y in range(-x - 1, 1)]) for x in range(3)), + (dict(**{str(k): k for k in range(x + 1)}) for x in range(3)), + ], +) +def test_star_argument_calls(genexpr): + """Test reconstruction of `*args`/`**kwargs` call sites (CALL_FUNCTION_EX). + + The arguments arrive already collected into a sequence and a mapping, so the + reconstruction spells every argument as unpacked; that evaluates identically + even where the source passed some of them plainly. + """ + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# LAMBDAS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ((lambda y: y * 2)(x) for x in range(5)), + ((lambda y, z: y + z)(x, x) for x in range(5)), + ((lambda y: y)(x) for x in range(5)), + # Closing over the loop variable, and over an enclosing comprehension + ((lambda y: y + x)(x) for x in range(5)), + ((lambda: x)() for x in range(5)), + (((lambda y: lambda z: z + y)(x))(1) for x in range(5)), + # Lambdas whose body is itself a comprehension + ((lambda: (y for y in range(x)))() for x in range(3)), + ((lambda: [y for y in range(x)])() for x in range(3)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(3)), + # Lambdas as arguments to other calls + (sorted([x, 1], key=lambda v: -v) for x in range(3)), + (list(map(lambda v: v * 2, [x, 1])) for x in range(3)), + # Conditional expressions inside a lambda body + ((lambda y: y if y % 2 else -y)(x) for x in range(5)), + ((lambda y: (y if y > 1 else 0) + 1)(x) for x in range(5)), + ], +) +def test_lambdas(genexpr): + """Test reconstruction of lambdas appearing inside comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Positional defaults, which attach to the *trailing* parameters + ((lambda y, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y=1: y)() for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2, w=3: y * z * w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=2: y * z)(x, 5) for x in range(3)), + # Keyword-only parameters, with and without defaults + ((lambda y, *, z=1: y + z)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z=1, w=2: y + z + w)(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, *, z: y + z)(x, z=4) for x in range(3)), + # Positional-only parameters + ((lambda y, /, z=2: y * z)(x) for x in range(3)), # type: ignore[assignment] + # *args and **kwargs + ((lambda *a: sum(a))(x, x) for x in range(3)), + ((lambda **k: sum(k.values()))(a=x) for x in range(3)), + ((lambda *a, **k: len(a) + len(k))(x, b=1) for x in range(3)), + ((lambda y, *a: y + len(a))(x, 1, 2) for x in range(3)), + ( + (lambda y, *a, z=3, **k: y + len(a) + z + len(k))(x, 1, w=2) + for x in range(3) + ), + # Defaults that are themselves non-trivial expressions + ((lambda y, z=(1, 2): y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ((lambda y, z=[1]: y + len(z))(x) for x in range(3)), # type: ignore[assignment] + ], +) +def test_lambda_default_and_variadic_arguments(genexpr): + """Test reconstruction of lambda defaults and variadic parameters.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# ASSIGNMENT EXPRESSIONS +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # A walrus binds in the *enclosing* scope, so at module level it + # compiles to COPY + STORE_GLOBAL rather than to a local. + (y for x in range(4) if (y := x * 2) > 1), # noqa: F821 + ((y := x) + 1 for x in range(3)), # noqa: F821 + (y * y for x in range(4) if (y := x + 1) > 2), # noqa: F821 + ((y := x) if x > 1 else -1 for x in range(4)), # noqa: F821 + # Bound in one clause and read in a later one + ((y, z) for x in range(3) if (y := x + 1) for z in range(y)), # noqa: F821 + # Inside a nested comprehension, and inside a lambda + ([(z := w) + z for w in range(x)] for x in range(4)), # noqa: F821 + ((lambda n: [(z := w) + z for w in range(n)])(x) for x in range(4)), # noqa: F821 + # Combined with a short-circuiting filter. The `or` must not re-evaluate + # the assignment, which is why the disjunction is absorbed. + (y for x in range(6) if (y := x * 2) > 6 or y == 0), # noqa: F821 + ], +) +def test_assignment_expressions(genexpr): + """A walrus in a comprehension binds in the *enclosing* scope (STORE_GLOBAL here).""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# COMPREHENSIONS NESTED IN EACH SYNTACTIC POSITION +# +# A comprehension can appear in the element, in the iterable, and inside a +# filter, and each of the four comprehension kinds can nest inside any other. +# The filter position is the interesting one: filters are reconstructed from +# control flow, so a comprehension inside a filter has to survive being treated +# as part of a boolean condition. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + # ... in filter position + (x for x in range(5) if any(y > 2 for y in range(x))), + (x for x in range(5) if all(y < 3 for y in range(x))), + (x for x in range(5) if [y for y in range(x)]), + (x for x in range(5) if {y for y in range(x)}), + (x for x in range(5) if {y: y for y in range(x)}), + (x for x in range(6) if len([y for y in range(x) if y % 2]) > 1), + (x for x in range(5) if any(y for y in range(x) if y % 2)), + ( + x + for x in range(6) + if all(y < x for y in range(2)) and any(z > 1 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in filter position, inside a short-circuiting condition + (x for x in range(6) if sum(y for y in range(x)) > 3 or x == 0), + (x for x in range(6) if x == 0 or any(y > 2 for y in range(x))), + (x for x in range(6) if x == 0 or len([y for y in range(x)]) > 2), + ( + x + for x in range(6) + if any(y > 1 for y in range(x)) or all(z < 2 for z in range(x)) + ), # noqa: E501 # fmt: skip + # ... in iterable position + (x for x in [y for y in range(5) if y % 2]), + (x for x in {y for y in range(5) if y % 2}), + (x for x in {y: y for y in range(3)}), + (x for x in (y for y in range(5) if y > 1 or y == 0)), + (x for x in [y for y in [z for z in range(4)] if y % 2]), + (x for x in [y for y in range(4)] if x > 1 or x == 0), + ((a, b) for a in range(3) for b in [c for c in range(a)]), + # ... in element position + ([y for y in range(x) if y % 2 or y == 0] for x in range(4)), + ({y for y in range(x) if y > 1} for x in range(5)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(4)), + (sum(y for y in range(x) if y % 2) for x in range(5)), + ([(z for z in range(y)) for y in range(x)] for x in range(3)), + # ... in several positions at once, with different kinds + ( + [y for y in {z for z in range(x)}] + for x in range(4) + if any(w > 1 for w in range(x)) + ), # noqa: E501 # fmt: skip + ({k: [v for v in range(k)] for k in {j for j in range(x)}} for x in range(4)), + ([y for y in range(x) if y or y == 0] for x in range(5) if x > 1 or x == 0), + ( + (y for y in range(x) if y % 2 or y == 0) + for x in (z for z in range(4)) + if x < 3 or x == 3 + ), # noqa: E501 # fmt: skip + # ... nesting the four kinds inside one another + ({k: {v for v in range(k)} for k in [j for j in range(x)]} for x in range(4)), + ([{y: y} for y in range(x)] for x in range(4)), + ({(y, y * 2) for y in range(x)} for x in range(4)), + ({y: y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ({y for y in range(x) if y % 2 or y == 0} for x in range(5)), + ([y for y in range(x) if 1 < y < 4] for x in range(6)), + ], +) +def test_comprehensions_in_every_position(genexpr): + """Test comprehensions nested in the element, the iterable and the filter.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda body is a separate code object with no loop of its own, so + # every branch inside one is a conditional expression rather than a + # filter -- including the branches of a comprehension nested in it. + ((lambda n: [y for y in range(n) if y % 2])(x) for x in range(4)), + ((lambda n: {y for y in range(n)})(x) for x in range(4)), + ((lambda n: {y: y**2 for y in range(n)})(x) for x in range(4)), + ((lambda n: sum(y for y in range(n)))(x) for x in range(4)), + ((lambda n: [y for y in range(n) if y > 1 or y == 0])(x) for x in range(5)), + ((lambda n: (y for y in range(n) if y % 2 or y == 0))(x) for x in range(4)), + ((lambda n: [y if y > 1 else -y for y in range(n)])(x) for x in range(4)), + ((lambda n: [y for y in range(n) if 1 < y < 3])(x) for x in range(5)), + # Lambdas nested in lambdas, and lambdas inside the comprehension body + ((lambda n: (lambda m: [y for y in range(m)])(n))(x) for x in range(3)), + ([(lambda v: v * 2)(y) for y in range(x)] for x in range(4)), + (list(map(lambda n: [y for y in range(n)], range(x))) for x in range(3)), + ([(lambda v: v if v > 1 else -v)(y) for y in range(x)] for x in range(4)), + ], +) +def test_comprehensions_inside_lambdas(genexpr): + """Test comprehensions nested inside lambda bodies.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +@pytest.mark.parametrize( + "genexpr", + [ + # Short-circuiting conditions nested in one another + (x for x in range(20) if (x > 2 or x < 1) and (x < 10 or x > 15)), + (x for x in range(20) if ((x > 2 and x < 5) or (x > 10 and x < 15)) or x == 0), + (x for x in range(30) if not (x % 2 == 0 or x % 3 == 0)), + (x for x in range(30) if not (not (x > 5) or not (x < 20))), + (x for x in range(40) if (x > 5 and x < 35) and (x % 3 == 0 or x % 5 == 0)), + ( + x + for x in range(40) + if (x < 5 and x % 2 == 0) or (10 < x < 15) or (x > 35 and x % 3 == 0) + ), # noqa: E501 # fmt: skip + # Short-circuiting conditions spanning several generators + ( + (x, y) + for x in range(6) + if x < 2 or x > 4 + for y in range(6) + if y < 1 or y > 4 + ), + ( + (x, y) + for x in range(5) + if x % 2 == 0 or x == 1 + for y in range(x) + if y > 0 and y < 3 + ), # noqa: E501 # fmt: skip + # Conditional expressions and filters that are both lazy + ((x if (x > 2 or x < 1) else -x) for x in range(10) if x % 2 == 0 or x == 1), + (x for x in range(20) if (x if x > 5 else not x) or x == 3), + ( + (x if x > 5 or x < 2 else (0 if x % 2 == 0 or x == 3 else 1)) + for x in range(12) + ), # noqa: E501 # fmt: skip + # Chained comparisons combined with lazy operators + (x for x in range(30) if 5 < x < 15 or 20 < x < 25), + (x for x in range(30) if 5 < x < 15 and (x % 2 == 0 or x % 3 == 0)), + ((x if 5 < x < 15 else 0) for x in range(20) if 2 < x < 18), + # Lazy conditions inside a nested comprehension, and around it + ([y for y in range(x) if y > 1 or y == 0] for x in range(5) if x > 2 or x == 0), + ((y for y in range(x) if y % 2 or y == 0) for x in range(4) if x < 3 or x == 3), + ], +) +def test_nested_lazy_conditions(genexpr): + """Test short-circuiting conditions nested inside one another.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# STRUCTURAL STRESS CASES +# ============================================================================ + +# These two must stay on one line: on Python 3.12 `dis` mis-reports jumps for +# multiline comprehensions, which test_multiline_comprehensions covers directly. +_STRESS_MANY_FILTERS = ((x, y) for x in range(10) if x % 2 == 0 if x > 2 for y in range(10) if y % 3 == 0 if y < x) # fmt: skip +_STRESS_NESTED_TERNARY = ([y if y > 1 else -y for y in range(x)] for x in range(4) if (x if x % 2 == 1 else x % 2 == 0)) # fmt: skip + + +@pytest.mark.parametrize( + "genexpr", + [ + # Deep loop nesting + ( + a + b + c + d + e + for a in range(2) + for b in range(2) + for c in range(2) + for d in range(2) + for e in range(2) + ), + ( + (a, b, c, d) + for a in range(2) + for b in range(a + 1) + for c in range(b + 1) + for d in range(c + 1) + ), + # Many filters spread over many loops. Kept on one line: on Python 3.12 + # `dis` mis-reports jumps for multiline comprehensions, which is covered + # separately by test_multiline_comprehensions below. + (x for x in range(50) if x > 5 if x < 40 if x % 2 == 0 if x % 3 == 0), + _STRESS_MANY_FILTERS, + # Deep comprehension nesting + (((z for z in range(y)) for y in range(x)) for x in range(3)), + ([[z for z in range(y)] for y in range(x)] for x in range(3)), + ({y: [z for z in range(y)] for y in range(x)} for x in range(3)), + # Structured literals in the element position + ({x, x + 1} for x in range(3)), + ({x: x + 1} for x in range(3)), + (((x, x), x) for x in range(3)), + ([x, [x, [x]]] for x in range(3)), + ({"k": [x, {"j": (x,)}]} for x in range(3)), + # Ternaries interleaved with nesting + _STRESS_NESTED_TERNARY, + (((y if y else -1) for y in range(x)) for x in range(3)), + ], +) +def test_structural_stress(genexpr): + """Test reconstruction of deeply nested and heavily filtered comprehensions.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +# ============================================================================ +# MULTILINE COMPREHENSIONS +# +# On Python 3.12 `dis` reports a different jump layout for a filter whose source +# spans several lines. The reconstruction used to come out negated, because the +# filter/conditional distinction was drawn from the *local* instruction order. +# Classifying branches from the control-flow graph instead is insensitive to +# that, so these now reconstruct identically on 3.12 and 3.13+. +# ============================================================================ + + +@pytest.mark.parametrize( + "genexpr", + [ + ( + x + for x in range(5) # comment to avoid reformatting + if x > 1 + ), + ( + x + for x in range(10) # comment to avoid reformatting + if x % 2 == 0 + if x > 2 + ), + ( + (x, y) + for x in range(10) + if x % 2 == 0 + if x > 2 + for y in range(10) + if y % 3 == 0 + if y < x + ), + ( + [y if y > 1 else -y for y in range(x)] + for x in range(4) + if (x if x % 2 == 1 else x % 2 == 0) + ), + ], +) +def test_multiline_comprehensions(genexpr): + """Filters in multiline comprehensions are mis-disassembled on Python 3.12.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + +def test_multiline_comprehensions_same_on_one_line(): + """The same expressions reconstruct correctly when written on one line.""" + one_line = (x for x in range(10) if x % 2 == 0 if x > 2) + assert_ast_equivalent(one_line, disassemble(one_line)) + + +# ============================================================================ +# HELPER FUNCTION TESTS +# ============================================================================ + + +@pytest.mark.parametrize( + "value,expected_str", + [ + # AST nodes should be returned as-is + (ast.Name(id="x", ctx=ast.Load()), "x"), + (ast.Constant(value=42), "42"), + (ast.List(elts=[], ctx=ast.Load()), "[]"), + ( + ast.BinOp( + left=ast.Constant(value=1), op=ast.Add(), right=ast.Constant(value=2) + ), + "1 + 2", + ), + # Constants should become ast.Constant nodes + (42, "42"), + (3.14, "3.14"), + (-42, "-42"), + (-3.14, "-3.14"), + ("hello", "'hello'"), + ("", "''"), + (b"bytes", "b'bytes'"), + (b"", "b''"), + (True, "True"), + (False, "False"), + (None, "None"), + # Complex numbers + (1 + 2j, "(1+2j)"), + (0 + 1j, "1j"), + (3 + 0j, "(3+0j)"), + (-1 - 2j, "(-1-2j)"), + # Tuples should become ast.Tuple nodes + ((), "()"), + ((1,), "(1,)"), + ((1, 2), "(1, 2)"), + (("a", "b", "c"), "('a', 'b', 'c')"), + # Special dict_item tuples + (("dict_item", "key", "value"), "('key', 'value')"), + (("dict_item", 42, "answer"), "(42, 'answer')"), + # Nested tuples + ((1, (2, 3)), "(1, (2, 3))"), + (((1, 2), (3, 4)), "((1, 2), (3, 4))"), + ((1, 2, (3, (4, 5))), "(1, 2, (3, (4, 5)))"), + # Lists should become ast.List nodes + ([1, 2, 3], "[1, 2, 3]"), + (["hello", "world"], "['hello', 'world']"), + ([True, False, None], "[True, False, None]"), + # Nested lists + ([[1, 2], [3, 4]], "[[1, 2], [3, 4]]"), + ([1, [2, [3, 4]], 5], "[1, [2, [3, 4]], 5]"), + # Mixed nested structures + ([(1, 2), (3, 4)], "[(1, 2), (3, 4)]"), + (([1, 2], [3, 4]), "([1, 2], [3, 4])"), + # Dicts should become ast.Dict nodes + ({"a": 1}, "{'a': 1}"), + ({"x": 10, "y": 20}, "{'x': 10, 'y': 20}"), + ({1: "one", 2: "two"}, "{1: 'one', 2: 'two'}"), + # Nested dicts + ({"a": {"b": 1}}, "{'a': {'b': 1}}"), + ( + {"nums": [1, 2, 3], "strs": ["a", "b"]}, + "{'nums': [1, 2, 3], 'strs': ['a', 'b']}", + ), + # Range objects + (range(5), "range(0, 5, 1)"), + (range(1, 10), "range(1, 10, 1)"), + (range(0, 10, 2), "range(0, 10, 2)"), + (range(10, 0, -1), "range(10, 0, -1)"), + (range(-5, 5), "range(-5, 5, 1)"), + # Empty collections + ([], "[]"), + ((), "()"), + ({}, "{}"), + # Complex nested structures + ([1, [2, 3], 4], "[1, [2, 3], 4]"), + ({"a": [1, 2], "b": {"c": 3}}, "{'a': [1, 2], 'b': {'c': 3}}"), + ([(1, {"a": [2, 3]}), ({"b": 4}, 5)], "[(1, {'a': [2, 3]}), ({'b': 4}, 5)]"), + # Edge cases with special values + ([None, True, False, 0, ""], "[None, True, False, 0, '']"), + ( + {"": "empty", None: "none", 0: "zero"}, + "{'': 'empty', None: 'none', 0: 'zero'}", + ), + # Large numbers + (999999999999999999999, "999999999999999999999"), + (1.7976931348623157e308, "1.7976931348623157e+308"), # Close to float max + # Sets - note unparse equivalence may fail for unordered collections + ({1, 2, 3}, "{1, 2, 3}"), + ], +) +def test_ensure_ast(value, expected_str): + """Test that ensure_ast correctly converts various values to AST nodes.""" + + result = ensure_ast(value) + + # Compare the unparsed strings + result_str = ast.unparse(result) + assert result_str == expected_str, ( + f"ensure_ast({repr(value)}) produced '{result_str}', expected '{expected_str}'" + ) + + +def test_error_handling(): + """Test that appropriate errors are raised for unsupported cases.""" + # Test with non-generator input + with pytest.raises(AssertionError): + disassemble([1, 2, 3]) # Not a generator + + # Test with consumed generator + gen = (x for x in range(5)) + list(gen) # Consume it + with pytest.raises(AssertionError): + disassemble(gen) + + +def test_comp_lambda_copy(): + """Test that CompLambda is compatible with copy.copy and copy.deepcopy.""" + # Create a test generator expression AST + genexpr_ast = ast.GeneratorExp( + elt=ast.Name(id="x", ctx=ast.Load()), + generators=[ + ast.comprehension( + target=ast.Name(id="x", ctx=ast.Store()), + iter=DummyIterName(), + ifs=[], + is_async=0, + ) + ], + ) + + # Create a CompLambda instance + comp_lambda = CompLambda(genexpr_ast) + + # Test copy.copy + copied = copy.copy(comp_lambda) + assert isinstance(copied, CompLambda) + assert ast.unparse(copied.body) == ast.unparse(comp_lambda.body) + assert copied.body is comp_lambda.body # Shallow copy shares the body + + # Test copy.deepcopy + deep_copied = copy.deepcopy(comp_lambda) + assert isinstance(deep_copied, CompLambda) + assert ast.unparse(deep_copied.body) == ast.unparse(comp_lambda.body) + assert deep_copied.body is not comp_lambda.body # Deep copy creates new body + + # Test that deep copied version works the same way + iterator = ast.Call( + func=ast.Name(id="range", ctx=ast.Load()), + args=[ast.Constant(value=5)], + keywords=[], + ) + + original_result = comp_lambda.inline(iterator) + deep_copied_result = deep_copied.inline(iterator) + + assert ast.unparse(original_result) == ast.unparse(deep_copied_result) + assert type(original_result) == type(deep_copied_result) + + +# ============================================================================ +# AST TRANSFORMER TESTS +# ============================================================================ From ccf7733066f380cf795dc92e85b55e8c9dd2f3f5 Mon Sep 17 00:00:00 2001 From: Eli Date: Tue, 28 Jul 2026 15:02:36 -0400 Subject: [PATCH 3/3] Address review comments on the generator-expression disassembler Six fixes, each with tests that fail without them: - `handle_build_map` read the key/value pairs of a dict display from the top of the stack down, reversing source order: a later duplicate key lost to an earlier one, and side effects ran backwards. - `_ensure_ast_tuple` treated any tuple whose first element was the string "dict_item" as an internal marker and dropped that element. Nothing produced such a marker; user data holding that string was silently corrupted. The special case is gone. - A free variable was reconstructed as a bare `ast.Name`, so evaluating the result resolved it against the evaluating namespace instead of the captured cell. The captured value is now written into the tree, for the generator itself, for lambdas reached as live objects, and for lambdas and comprehensions nested inside. A cell the comprehension creates -- a target captured by a nested lambda -- still stands as a name, since the reconstruction binds it too. A capture with no AST spelling, including an iterator, raises `TypeError` rather than reconstructing to a name that would answer differently. - `_ensure_ast_iterator_adaptor` ignored the strictness a `zip` pickles as reduction state, so a strict zip silently truncated ragged input where the original raised. - A lambda reached as a live object lost its default values, which live on the function rather than in its code object, leaving parameters with no way to be filled. - `disassemble` asserted on its input; it now raises `ValueError`, and checks the generator has not been started rather than leaving that to an assert further in. Also documents what reconstruction does and does not recover: evaluating the result re-runs every expression in it, so a stateful filter answers against state as it then stands. 663 passed, 2 xfailed on 3.12, 3.13 and 3.14. --- effectful/internals/disassembly.py | 214 +++++++++++++--- tests/test_internals_disassembler.py | 355 ++++++++++++++++++++++++++- 2 files changed, 533 insertions(+), 36 deletions(-) diff --git a/effectful/internals/disassembly.py b/effectful/internals/disassembly.py index 47aaccb52..d2132b193 100644 --- a/effectful/internals/disassembly.py +++ b/effectful/internals/disassembly.py @@ -13,6 +13,33 @@ >>> g = (x * 2 for x in range(10) if x % 2 == 0) >>> ast_node = disassemble(g) >>> # ast_node is now an ast.Expression representing the original expression + +What is recovered, and what is not: + + What comes back is the comprehension's *syntax*, not the generator's + suspended state, so evaluating the reconstruction runs every expression in + it a second time. For a comprehension over pure expressions the two agree + element for element. Where they can part company: + + * A **stateful filter or element expression** -- one that mutates something, + or whose value depends on state that has moved on -- is re-run, and can + answer differently the second time. `(x for x in xs if next(flags))` + reconstructs faithfully as source, but iterating the reconstruction draws + from `flags` where it now stands, not from where it stood when the + original generator was built. The side effects happen twice, once per + iteration of each generator. + + * The **outermost iterable** is not part of the comprehension's bytecode -- + it is an object the generator already holds -- so it is recovered from + that object rather than from source. What lands in the tree is a snapshot + of the elements not yet consumed: `iter([1, 2, 3])` advanced once becomes + the literal `[2, 3]`. The expression that produced it is gone, and so is + any laziness it had. + + * A **free variable** is likewise recovered by value, not by name: the value + in the closure cell is written into the tree. A captured object with no + AST spelling raises `TypeError` rather than reconstructing to a name that + would resolve against the evaluating namespace instead. """ import ast @@ -318,6 +345,12 @@ class ReconstructionState: # KW_NAMES has no stack effect, so the names cannot live on `stack`. kw_names: tuple[str, ...] | None = field(default=None) + # The value captured in each closure cell the code reads, by name. A free + # variable is looked up in a cell, not in globals, so reconstructing it as a + # bare name would resolve to whatever the evaluating namespace happens to + # bind; the captured value has to be written into the tree instead. + freevars: dict[str, ast.expr] = field(default_factory=dict) + @property def instructions(self) -> collections.abc.Mapping[int, dis.Instruction]: """The bytecode instructions of the current code object, by offset.""" @@ -959,7 +992,7 @@ def _decide_branch( return replace(state, branches={**state.branches, instr.offset: edge}) -def _symbolic_exec(code: types.CodeType) -> ast.expr: +def _symbolic_exec(code: types.CodeType, freevars: dict[str, ast.expr]) -> ast.expr: """Execute bytecode symbolically, following control flow.""" continuations: list[ReconstructionState] = [ ReconstructionState( @@ -969,6 +1002,7 @@ def _symbolic_exec(code: types.CodeType) -> ast.expr: if current_version() == PythonVersion.PY_312 and code.co_flags & inspect.CO_GENERATOR else [Placeholder()], + freevars=freevars, ) ] @@ -1168,12 +1202,14 @@ def handle_build_map( new_stack = state.stack + [new_result] return replace(state, stack=new_stack) else: - # Pop key-value pairs for the dict - keys: list[ast.expr | None] = [ - ensure_ast(state.stack[-2 * i - 2]) for i in range(size) - ] - values = [ensure_ast(state.stack[-2 * i - 1]) for i in range(size)] - new_stack = state.stack[: -2 * size] if size > 0 else state.stack + # Pop key-value pairs for the dict. They sit on the stack in source + # order -- key first, then value -- so they must be read from the + # bottom of that slice up: a later duplicate key has to keep winning, + # and side effects have to happen in the order they were written. + pairs = state.stack[-2 * size :] + keys: list[ast.expr | None] = [ensure_ast(pairs[2 * i]) for i in range(size)] + values = [ensure_ast(pairs[2 * i + 1]) for i in range(size)] + new_stack = state.stack[: -2 * size] # Create dict AST dict_node = ast.Dict(keys=keys, values=values) @@ -1453,10 +1489,18 @@ def handle_load_fast( def handle_load_deref( state: ReconstructionState, instr: dis.Instruction ) -> ReconstructionState: - # LOAD_DEREF loads a value from a closure variable + # LOAD_DEREF loads a value out of a cell. When the cell is one this code + # object *creates* (a comprehension target captured by a nested lambda, say) + # the name is bound inside the reconstructed tree too, so it can stand. A + # free variable is different: its cell belongs to an enclosing scope that + # the reconstructed tree does not reproduce, so a bare name would silently + # become a global lookup. Write the captured value in instead. var_name = instr.argval - new_stack = state.stack + [ast.Name(id=var_name, ctx=ast.Load())] - return replace(state, stack=new_stack) + if var_name in state.code.co_freevars and var_name in state.freevars: + loaded: ast.expr = copy.deepcopy(state.freevars[var_name]) + else: + loaded = ast.Name(id=var_name, ctx=ast.Load()) + return replace(state, stack=state.stack + [loaded]) @register_handler("LOAD_CLOSURE", version=PythonVersion.PY_312) @@ -1478,8 +1522,15 @@ def handle_load_const( state: ReconstructionState, instr: dis.Instruction ) -> ReconstructionState: const_value = instr.argval - new_stack = state.stack + [ensure_ast(const_value)] - return replace(state, stack=new_stack) + # A nested lambda or comprehension arrives as a code object, and its own + # free variables reach back through this scope to the same cells, so the + # captured values have to travel with it. + loaded = ( + _reconstruct_code(const_value, state.freevars) + if isinstance(const_value, types.CodeType) + else ensure_ast(const_value) + ) + return replace(state, stack=state.stack + [loaded]) @register_handler("LOAD_GLOBAL", version=PythonVersion.PY_312) @@ -2201,7 +2252,9 @@ def _apply_function_attribute( ) -> ast.Lambda | CompLambda: """Attach one function attribute to a reconstructed lambda.""" if flag == MAKE_FUNCTION_CLOSURE: - # Free variables are already spelled by name in the reconstructed body. + # The body has already resolved each free variable: to a captured value + # if the cell came from outside the comprehension, and otherwise to the + # name the reconstructed tree binds it under. See `handle_load_deref`. return func if flag == MAKE_FUNCTION_ANNOTATE: # A lambda has no annotations, and the AST does not carry the lazy @@ -3108,13 +3161,8 @@ def _ensure_ast_constant(value) -> ast.Constant: @ensure_ast.register def _ensure_ast_tuple(value: tuple) -> ast.Tuple: - """Convert tuple to AST - special handling for dict items""" - if len(value) > 0 and value[0] == "dict_item": - return ast.Tuple( - elts=[ensure_ast(value[1]), ensure_ast(value[2])], ctx=ast.Load() - ) - else: - return ast.Tuple(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) + """Convert tuple to AST""" + return ast.Tuple(elts=[ensure_ast(v) for v in value], ctx=ast.Load()) def _unconsumed(value: Iterator) -> typing.Any: @@ -3219,15 +3267,22 @@ def _ensure_ast_iterator_adaptor(value: Iterator) -> ast.Call: or range iterator they cannot be materialised -- but ``__reduce__`` hands back their constituent parts, each of which ``ensure_ast`` can handle in turn. Any already-consumed prefix is reflected in the inner iterators. + + A ``zip`` also pickles its strictness as reduction state, which has to be + carried over as a keyword argument: a strict ``zip`` raises on ragged input + where a lax one stops at the shortest iterable. """ reduced = value.__reduce__() if isinstance(reduced, str): raise TypeError(f"Cannot convert {type(value)} to AST node") func, args = reduced[:2] + keywords = [] + if isinstance(value, zip) and len(reduced) > 2 and reduced[2]: + keywords.append(ast.keyword(arg="strict", value=ast.Constant(value=True))) return ast.Call( func=ast.Name(id=func.__name__, ctx=ast.Load()), args=[ensure_ast(arg) for arg in args], - keywords=[], + keywords=keywords, ) @@ -3255,8 +3310,62 @@ def _ensure_ast_range_iterator(value: Iterator) -> ast.Call: return ensure_ast(_unconsumed(value)) # type: ignore -@ensure_ast.register -def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: +def _cell_values( + code: types.CodeType, closure: tuple[types.CellType, ...] | None +) -> dict[str, typing.Any]: + """Read the cells a function closed over, by free-variable name. + + A cell that is still empty -- a recursive definition not yet bound, say -- + is left out. + """ + values: dict[str, typing.Any] = {} + for name, cell in zip(code.co_freevars, closure or ()): + try: + values[name] = cell.cell_contents + except ValueError: + continue + return values + + +def _freevar_bindings( + code: types.CodeType, captured: collections.abc.Mapping[str, typing.Any] +) -> dict[str, ast.expr]: + """AST nodes for the values ``code``'s free variables were captured from. + + A free variable with no captured value to be found is left out, and the + reconstructed tree keeps the bare name, which is as close as it can get. + + An iterator is refused rather than written in. `ensure_ast` spells one as + the elements it has left, which is what the outermost iterable wants -- it + is about to be consumed anyway -- but a captured iterator is a value the + body can do anything with, and a list of its remaining elements is not the + same object. + """ + bindings: dict[str, ast.expr] = {} + for name in code.co_freevars: + if name not in captured: + continue + value = captured[name] + try: + if isinstance(value, Iterator): + raise TypeError("an iterator has no AST spelling in value position") + bindings[name] = ensure_ast(value) + except (TypeError, AssertionError) as exc: + raise TypeError( + f"Cannot represent {value!r}, " + f"the value captured in free variable {name!r}: {exc}" + ) from exc + return bindings + + +def _reconstruct_code( + value: types.CodeType, freevars: dict[str, ast.expr] +) -> ast.Lambda | CompLambda: + """Reconstruct a lambda or comprehension body from its code object. + + ``freevars`` maps the names this code captures from enclosing scopes to the + values found in their cells; see `handle_load_deref`. + """ assert inspect.iscode(value), "Input must be a code object" name: str = value.co_name.split(".")[-1] @@ -3277,7 +3386,7 @@ def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: raise TypeError(f"Unsupported code object type: {value.co_name}") # Symbolic execution to reconstruct the AST - result: ast.expr = _symbolic_exec(value) + result: ast.expr = _symbolic_exec(value, freevars) # Check postconditions assert not any(isinstance(x, ast.stmt) for x in ast.walk(result)), ( @@ -3349,15 +3458,38 @@ def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: raise TypeError(f"Invalid result for type {name}: {result}") +@ensure_ast.register +def _ensure_ast_codeobj(value: types.CodeType) -> ast.Lambda | CompLambda: + """A bare code object has no cells attached, so free variables stay names.""" + return _reconstruct_code(value, {}) + + @ensure_ast.register def _ensure_ast_lambda(value: types.LambdaType) -> ast.Lambda: assert inspect.isfunction(value) and value.__name__.endswith(""), ( "Input must be a lambda function" ) code: types.CodeType = value.__code__ - result = ensure_ast(code) + result = _reconstruct_code( + code, _freevar_bindings(code, _cell_values(code, value.__closure__)) + ) assert isinstance(result, ast.Lambda), "Lambda body must be an AST Lambda node" assert not isinstance(result, CompLambda), "Lambda must not be a CompLambda" + + # Default values are not in the code object: they were evaluated where the + # lambda was written and attached to the function. A lambda built inside the + # comprehension gets them from the stack instead -- see + # `_apply_function_attribute` -- but one arriving as a live object carries + # them here, and dropping them would leave parameters with no way to be + # filled. They cover the *trailing* positional parameters. + if value.__defaults__: + result.args.defaults = [ensure_ast(d) for d in value.__defaults__] + if value.__kwdefaults__: + by_name = value.__kwdefaults__ + result.args.kw_defaults = [ + ensure_ast(by_name[arg.arg]) if arg.arg in by_name else None + for arg in result.args.kwonlyargs + ] return result @@ -3367,10 +3499,15 @@ def _ensure_ast_genexpr(genexpr: types.GeneratorType) -> ast.GeneratorExp: assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( "Generator must be in created state" ) - genexpr_ast = ensure_ast(genexpr.gi_code) - assert isinstance(genexpr_ast, CompLambda) assert genexpr.gi_frame is not None, "Generator must not be exhausted" - geniter_ast = ensure_ast(genexpr.gi_frame.f_locals[".0"]) + # A generator holds its cells in its frame rather than in a __closure__, and + # an unstarted frame already has them all copied in. + frame_locals = genexpr.gi_frame.f_locals + genexpr_ast = _reconstruct_code( + genexpr.gi_code, _freevar_bindings(genexpr.gi_code, frame_locals) + ) + assert isinstance(genexpr_ast, CompLambda) + geniter_ast = ensure_ast(frame_locals[".0"]) result = genexpr_ast.inline(geniter_ast) assert isinstance(result, ast.GeneratorExp) assert inspect.getgeneratorstate(genexpr) == inspect.GEN_CREATED, ( @@ -3411,8 +3548,11 @@ def disassemble( ast.Expression: An AST node representing the reconstructed comprehension. Raises: - AssertionError: If the input is not a generator or if the generator + ValueError: If the input is not a generator or if the generator has already been started (not in 'GEN_CREATED' state). + TypeError: If some part of the comprehension has no AST spelling -- + an outermost iterable or a captured free variable whose value + cannot be written into the tree. Example: >>> # Generator expression @@ -3432,7 +3572,19 @@ def disassemble( The reconstruction is based on bytecode analysis and may not perfectly preserve the original source code formatting or variable names in all cases. However, the semantic behavior of the reconstructed AST should - match the original comprehension. + match the original comprehension, subject to the limits set out under + "What is recovered, and what is not" in this module's docstring: + evaluating the result re-runs every expression in it, so a stateful + filter or element expression can answer differently the second time, + and the outermost iterable is a snapshot rather than an expression. """ - assert inspect.isgenerator(genexpr), "Input must be a generator expression" + if not inspect.isgenerator(genexpr): + raise ValueError( + f"Input must be a generator expression, got {type(genexpr).__name__}" + ) + if inspect.getgeneratorstate(genexpr) != inspect.GEN_CREATED: + raise ValueError( + "Input must be a generator expression that has not been started, " + f"got one in state {inspect.getgeneratorstate(genexpr)}" + ) return ast.fix_missing_locations(ast.Expression(ensure_ast(genexpr))) diff --git a/tests/test_internals_disassembler.py b/tests/test_internals_disassembler.py index 55d947f6e..22e484167 100644 --- a/tests/test_internals_disassembler.py +++ b/tests/test_internals_disassembler.py @@ -519,6 +519,63 @@ def test_different_comprehension_types(genexpr): assert_ast_equivalent(genexpr, ast_node) +# ============================================================================ +# DICT DISPLAYS +# +# A dict display with dynamic keys is built by BUILD_MAP from key/value pairs +# that the compiler pushed in source order. Dicts compare equal whatever order +# they were built in, so these tests pin the order down directly: which of two +# equal keys wins, what `items()` yields, and when each subexpression runs. +# ============================================================================ + + +_EVAL_ORDER: list[str] = [] + + +def _note(tag: str, value: typing.Any) -> typing.Any: + """Record that this subexpression was evaluated, and pass its value along.""" + _EVAL_ORDER.append(tag) + return value + + +def test_dict_display_duplicate_keys(): + """The last of two equal keys wins, so the pairs must keep their order.""" + genexpr = ({x: "first", x: "second"} for x in range(1)) # noqa: F602 + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "({x: 'first', x: 'second'} for x in range(0, 1, 1))" + ) + assert materialize(genexpr) == [{0: "second"}] + assert materialize(compile_and_eval(reconstructed)) == [{0: "second"}] + + +def test_dict_display_insertion_order(): + genexpr = ({x: "a", x + 1: "b", x + 2: "c"} for x in range(1)) + reconstructed = disassemble(genexpr) + expected = [[(0, "a"), (1, "b"), (2, "c")]] + assert [list(d.items()) for d in genexpr] == expected + assert [list(d.items()) for d in compile_and_eval(reconstructed)] == expected + + +def test_dict_display_evaluation_order(): + """Keys and values run left to right, key before its own value.""" + genexpr = ( + {_note("k1", "a"): _note("v1", 1), _note("k2", "b"): _note("v2", 2)} + for _ in range(1) + ) + reconstructed = disassemble(genexpr) + + _EVAL_ORDER.clear() + assert materialize(genexpr) == [{"a": 1, "b": 2}] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + _EVAL_ORDER.clear() + assert materialize(compile_and_eval(reconstructed, {"_note": _note})) == [ + {"a": 1, "b": 2} + ] + assert _EVAL_ORDER == ["k1", "v1", "k2", "v2"] + + # ============================================================================ # CONDITIONAL EXPRESSIONS # ============================================================================ @@ -1006,6 +1063,11 @@ def test_outermost_iterable_types(genexpr): (x for x in map(max, [1, 2], [3, 0])), (x for x in filter(None, [0, 1, 2])), (x for x in filter(bool, [0, 1, 2])), + # A strict zip over equal-length iterables behaves like a lax one, but + # its strictness has to survive anyway -- see the ragged case below. + ((a, b) for a, b in zip([1, 2], [3, 4], strict=True)), + (x for x in zip("ab", range(2), [7, 8], strict=True)), + (x for x in zip(range(2), map(abs, [-1, -2]), strict=True)), # Nested adaptors, and adaptors over non-sequence iterables (x for x in zip(range(2), map(abs, [-1, -2]))), (x for x in enumerate(filter(None, [0, 1]))), @@ -1086,6 +1148,59 @@ def test_unreachable_comprehension_body(genexpr): assert materialize(compile_and_eval(ast_node)) == [] +def test_outermost_iterable_strict_zip_stays_strict(): + """A ragged strict zip raises; the reconstruction must raise there too.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5], strict=True)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == ( + "(a + b for a, b in zip([1, 2, 3], [4, 5], strict=True))" + ) + + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + with pytest.raises(ValueError): + materialize(genexpr) + + +def test_outermost_iterable_lax_zip_stays_lax(): + """A lax zip stops at the shortest iterable and must not become strict.""" + genexpr = (a + b for a, b in zip([1, 2, 3], [4, 5])) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(a + b for a, b in zip([1, 2, 3], [4, 5]))" + assert materialize(compile_and_eval(reconstructed)) == [5, 7] + assert materialize(genexpr) == [5, 7] + + +def test_outermost_iterable_partially_consumed_strict_zip(): + """Strictness survives even once the zip has been partly consumed.""" + zipped = zip([1, 2, 3], [4, 5], strict=True) + next(zipped) + + genexpr = (a + b for a, b in zipped) + reconstructed = disassemble(genexpr) # must precede consuming `genexpr` + assert "strict=True" in ast.unparse(reconstructed) + with pytest.raises(ValueError): + materialize(compile_and_eval(reconstructed)) + + +@pytest.mark.parametrize( + "genexpr", + [ + # "dict_item" was once an internal marker in the first slot of a tuple, + # which is a string a user's data is perfectly entitled to hold. + (x for x in (("dict_item", 1, 2),)), + (x for x in [("dict_item", "key", "value")]), + (x for x in (("dict_item",), ("dict_item", 1), ("dict_item", 1, 2, 3))), + (("dict_item", x) for x in range(2)), + (("dict_item", x, x + 1) for x in range(2)), + (x for x in {("dict_item", 1, 2): "v"}.items()), + ], +) +def test_tuples_starting_with_dict_item(genexpr): + """No element of a user's tuple is an internal marker to be stripped.""" + assert_ast_equivalent(genexpr, disassemble(genexpr)) + + def test_outermost_iterable_partially_consumed_adaptor(): """A consumed prefix is reflected in the adaptor's inner iterators.""" zipped = zip([1, 2, 3], [4, 5, 6]) @@ -1302,6 +1417,32 @@ def test_lambda_default_and_variadic_arguments(genexpr): assert_ast_equivalent(genexpr, ast_node) +@pytest.mark.parametrize( + "genexpr", + [ + # A lambda reached as a live object -- through the outermost iterable -- + # rather than built inside the comprehension. Its defaults are on the + # function, not in the code object it was compiled to. + (fn() for fn in [lambda value=1: value]), + (fn(2) for fn in [lambda a, b=10: a * b]), + (fn(2) for fn in (lambda a, b=10: a * b,)), + (fn(1) for fn in [lambda a, *, k=5: a + k]), + (fn(1) for fn in [lambda a, /, b=2, *, k=5: a * b + k]), + (fn() for fn in [lambda x=(1, 2): sum(x)]), + (fn() for fn in [lambda x=[1, 2]: len(x)]), + (fn(1, 2) for fn in [lambda a, b=0, *rest, k=5, **kw: a + b + k + len(rest)]), + # ... several of them, and one with no defaults alongside + (fn() for fn in [lambda: 0, lambda v=1: v, lambda v=2: v]), # type: ignore[misc] + # ... and one passed through map rather than iterated directly + (fn(3) for fn in map(lambda f: f, [lambda a, b=4: a * b])), + ], +) +def test_lambda_object_defaults(genexpr): + """A lambda arriving as a live object keeps the defaults attached to it.""" + ast_node = disassemble(genexpr) + assert_ast_equivalent(genexpr, ast_node) + + # ============================================================================ # ASSIGNMENT EXPRESSIONS # ============================================================================ @@ -1437,6 +1578,204 @@ def test_comprehensions_inside_lambdas(genexpr): assert_ast_equivalent(genexpr, ast_node) +# ============================================================================ +# CLOSURES +# +# A comprehension written inside a function reads the function's locals out of +# closure cells, not out of globals. Those cells belong to a scope the +# reconstruction does not reproduce, so leaving a free variable as a bare name +# would silently turn it into a global lookup -- picking up a different value, +# or none at all. The captured value is written into the tree instead. Each +# test below evaluates the reconstruction in a namespace that binds the same +# name to something else, so a lookup that leaked out would be caught. +# ============================================================================ + + +def test_closure_shadowed_by_global(): + def make(): + value = 1 + return (value for _ in range(1)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(1 for _ in range(0, 1, 1))" + assert materialize(genexpr) == [1] + assert materialize(compile_and_eval(reconstructed, {"value": 2})) == [1] + + +@pytest.mark.parametrize( + "make,shadow,expected", + [ + # In the element expression, the filter, and an inner iterable + (lambda: (lambda n: (x * n for x in range(4)))(3), {"n": 100}, [0, 3, 6, 9]), + ( + lambda: (lambda t: (x for x in range(6) if x > t))(3), + {"t": -1}, + [4, 5], + ), + ( + lambda: (lambda k: (y for x in range(2) for y in range(k)))(2), + {"k": 5}, + [0, 1, 0, 1], + ), + # A captured container, indexed and iterated + ( + lambda: (lambda d: (d[i] for i in range(2)))([10, 20]), + {"d": [0, 0]}, + [10, 20], + ), + ( + lambda: (lambda d: (v for v in d))({"a": 1, "b": 2}), + {"d": {}}, + ["a", "b"], + ), + # Captured by a lambda nested inside the comprehension + ( + lambda: (lambda n: ((lambda y: y * n)(x) for x in range(4)))(3), + {"n": 100}, + [0, 3, 6, 9], + ), + ( + lambda: (lambda n: ((lambda: (lambda: n)())() for _ in range(2)))(7), + {"n": 0}, + [7, 7], + ), + # Captured by a comprehension nested inside the comprehension + ( + lambda: (lambda n: ([y * n for y in range(x)] for x in range(3)))(2), + {"n": 100}, + [[], [0], [0, 2]], + ), + ( + lambda: (lambda n: (sum(y for y in range(x) if y < n) for x in range(4)))( + 2 + ), + {"n": 100}, + [0, 0, 1, 1], + ), + # Two free variables at once + ( + lambda: (lambda a, b: (x * a + b for x in range(3)))(2, 1), + {"a": 0, "b": 0}, + [1, 3, 5], + ), + ], +) +def test_closure_free_variables(make, shadow, expected): + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == expected + assert materialize(compile_and_eval(reconstructed, dict(shadow))) == expected + # The name must not survive anywhere in the tree, or the shadowing binding + # above would have been the one that answered. + assert not ( + {node.id for node in ast.walk(reconstructed) if isinstance(node, ast.Name)} + & set(shadow) + ) + + +def test_closure_target_captured_by_nested_lambda_stays_a_name(): + """A cell this comprehension *creates* is bound in the reconstruction too.""" + genexpr = ((lambda: x)() for x in range(3)) + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [0, 1, 2] + assert materialize(compile_and_eval(reconstructed, {"x": 99})) == [0, 1, 2] + + +def test_closure_shadowed_by_a_nested_comprehension_target(): + """Only the free `n` is the captured one; the inner comprehension rebinds it.""" + + def make(): + n = 5 + return (n + sum(n for n in range(x)) for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 5, 6] + assert materialize(compile_and_eval(reconstructed, {"n": 100, "sum": sum})) == [ + 5, + 5, + 6, + ] + + +def test_closure_shadowed_by_a_nested_lambda_parameter(): + """Only the free `n` is the captured one; the lambda's parameter is its own.""" + + def make(): + n = 5 + return ((lambda n: n * 2)(x) + n for x in range(3)) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [5, 7, 9] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [5, 7, 9] + + +def test_closure_inside_a_live_lambda(): + """A lambda reached as an object closes over cells of its own.""" + + def make(): + n = 3 + return (fn(2) for fn in [lambda a, b=10: a * n + b]) + + genexpr = make() + reconstructed = disassemble(genexpr) + assert materialize(genexpr) == [16] + assert materialize(compile_and_eval(reconstructed, {"n": 100})) == [16] + + +def test_closure_value_that_cannot_be_represented(): + """A capture with no AST spelling is refused rather than quietly dropped.""" + + class Opaque: + pass + + def make(): + obj = Opaque() + return (obj for _ in range(1)) + + with pytest.raises(TypeError, match="captured in free variable 'obj'"): + disassemble(make()) + + +def test_closure_captured_iterator_is_refused(): + """An iterator's remaining elements are not the iterator, so it is refused.""" + + def make(): + flags = iter([True, False, True, False]) + return (x for x in range(4) if next(flags)) + + with pytest.raises(TypeError, match="captured in free variable 'flags'"): + disassemble(make()) + + +# ============================================================================ +# STATEFUL EXPRESSIONS +# +# What comes back is syntax, so evaluating it runs the comprehension a second +# time. A filter or element expression that depends on state that has since +# moved on answers differently then -- faithfully reconstructed, but no longer +# in agreement with the generator it came from. +# ============================================================================ + + +_FLAGS = iter([True, False, True, False, False, False, True, False]) + + +def test_stateful_filter_is_re_evaluated(): + genexpr = (x for x in range(4) if next(_FLAGS)) + reconstructed = disassemble(genexpr) + assert ast.unparse(reconstructed) == "(x for x in range(0, 4, 1) if next(_FLAGS))" + + # The reconstruction is the same comprehension, but `_FLAGS` has advanced by + # the time it runs, so the two do not agree element for element. + assert materialize(genexpr) == [0, 2] + assert materialize( + compile_and_eval(reconstructed, {"_FLAGS": _FLAGS, "next": next}) + ) == [2] + + @pytest.mark.parametrize( "genexpr", [ @@ -1636,9 +1975,9 @@ def test_multiline_comprehensions_same_on_one_line(): ((1,), "(1,)"), ((1, 2), "(1, 2)"), (("a", "b", "c"), "('a', 'b', 'c')"), - # Special dict_item tuples - (("dict_item", "key", "value"), "('key', 'value')"), - (("dict_item", 42, "answer"), "(42, 'answer')"), + # A tuple is a tuple: no element of one is a marker to be stripped + (("dict_item", "key", "value"), "('dict_item', 'key', 'value')"), + (("dict_item", 42, "answer"), "('dict_item', 42, 'answer')"), # Nested tuples ((1, (2, 3)), "(1, (2, 3))"), (((1, 2), (3, 4)), "((1, 2), (3, 4))"), @@ -1705,13 +2044,19 @@ def test_ensure_ast(value, expected_str): def test_error_handling(): """Test that appropriate errors are raised for unsupported cases.""" # Test with non-generator input - with pytest.raises(AssertionError): + with pytest.raises(ValueError): disassemble([1, 2, 3]) # Not a generator # Test with consumed generator gen = (x for x in range(5)) list(gen) # Consume it - with pytest.raises(AssertionError): + with pytest.raises(ValueError): + disassemble(gen) + + # Test with a generator that has been started but not consumed + gen = (x for x in range(5)) + next(gen) + with pytest.raises(ValueError): disassemble(gen)