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/effectful/internals/disassembly.py b/effectful/internals/disassembly.py new file mode 100644 index 000000000..d2132b193 --- /dev/null +++ b/effectful/internals/disassembly.py @@ -0,0 +1,3590 @@ +""" +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 + +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 +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) + + # 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.""" + 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, freevars: dict[str, ast.expr]) -> 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()], + freevars=freevars, + ) + ] + + 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. 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) + 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 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 + 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) +@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 + # 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) +@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: + # 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 + # 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""" + 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. + + 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, + ) + + +@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 + + +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] + + # 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, freevars) + + # 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_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 = _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 + + +@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" + ) + assert genexpr.gi_frame is not None, "Generator must not be exhausted" + # 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, ( + "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: + 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 + >>> 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, 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. + """ + 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_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 diff --git a/tests/test_internals_disassembler.py b/tests/test_internals_disassembler.py new file mode 100644 index 000000000..22e484167 --- /dev/null +++ b/tests/test_internals_disassembler.py @@ -0,0 +1,2109 @@ +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) + + +# ============================================================================ +# 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 +# ============================================================================ + + +@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])), + # 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]))), + (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_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]) + 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) + + +@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 +# ============================================================================ + + +@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) + + +# ============================================================================ +# 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", + [ + # 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')"), + # 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))"), + ((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(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(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) + + +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 +# ============================================================================