# -*- coding: utf-8 -*- """Multi-shot generators. A `@multishot` function is a generator-shaped construct whose execution state is captured *as a continuation* at every `myield`, so it can be resumed from any earlier `myield` arbitrarily many times — branching execution into independent timelines. Built on top of `call_cc[]` / `get_cc()`. Only meaningful inside a `with continuations:` block (this is enforced by `call_cc[]`, which the expansion of `myield` produces). For attribution: the one-shot vs. multi-shot continuation distinction goes back at least to Bruggeman, Waddell & Dybvig 1996 ("Representing control in the presence of one-shot continuations"). Racket's docs are the canonical reference for current usage: https://docs.racket-lang.org/reference/cont.html Public surface: - `multishot` — decorator macro that turns a `def` into a multi-shot generator. Inside, the four `myield` forms are recognized and rewritten. - `myield` — name/expr macro for the four yield variants: Multi-shot yield Returns `k` expects Single-shot analog myield k no argument yield myield[expr] (k, value) no argument yield expr var = myield k one argument var = yield var = myield[expr] (k, value) one argument var = yield expr - `MultishotIterator` — adapter that makes a `@multishot` conform to a subset of Python's generator protocol, plus `copy.copy()` for forking. See `doc/macros.md` for the user-facing documentation. """ import ast from functools import partial from mcpyrate.quotes import macros, q, n, a, h # noqa: F401 from mcpyrate import namemacro, gensym from mcpyrate.quotes import is_captured_value from mcpyrate.utils import extract_bindings, rename from mcpyrate.walkers import ASTTransformer from ..fun import identity from ..misc import safeissubclass from .scopeanalyzer import isnewscope from .tailtools import macros, call_cc # noqa: F401, F811 -- macro-import: makes `h[call_cc]` a hygienic *macro* reference from .tailtools import get_cc, iscontinuation __all__ = ["multishot", "myield", "myield_from", "MultishotIterator"] # -------------------------------------------------------------------------------- # `myield` — name/expr macro def myield_function(tree, syntax, **kw): """[syntax, name/expr] Yield from a multi-shot generator. Only meaningful at the top level of a function decorated with `@multishot`. Outside that context, raises `SyntaxError` at macro-expansion time. For details, see `multishot`. """ if syntax not in ("name", "expr"): raise SyntaxError("myield is a name and expr macro only") # pragma: no cover # Allow `myield` in non-Load contexts so the name can be assigned to / del'd # without spuriously triggering the macro (mostly defensive — `multishot` # itself recognizes the patterns it needs before this macro runs). if type(getattr(tree, "ctx", None)) in (ast.Store, ast.Del): return tree # `myield` is not really a macro; it's a marker that `@multishot` looks for # and rewrites away. If a `myield` survives to reach the expander, it was # placed somewhere `@multishot` couldn't see it (outside `@multishot`, or # inside a nested scope that `@multishot` deliberately doesn't recurse into). raise SyntaxError("myield may only appear at the top level of a `@multishot` generator") myield = namemacro(myield_function) # -------------------------------------------------------------------------------- # `myield_from` — name/expr macro def myield_from_function(tree, syntax, **kw): """[syntax, name/expr] Delegate to another `@multishot` generator. Multi-shot analog of `yield from`. Drives the inner multi-shot generator, re-yielding each of its values to the outer's caller. On inner exhaustion (`StopIteration`), execution continues in the outer body. The two forms:: myield_from[expr] # statement; inner's StopIteration value discarded var = myield_from[expr] # statement; inner's StopIteration value bound to var Only meaningful at the top level of a `@multishot` function decorated within a `with continuations:` block. Forwards `send` and `throw` from the outer's caller into the inner; while delegating, `outer_mi.gi_yieldfrom` returns the inner `MultishotIterator`. For details, see `multishot`. """ if syntax not in ("name", "expr"): raise SyntaxError("myield_from is a name and expr macro only") # pragma: no cover if type(getattr(tree, "ctx", None)) in (ast.Store, ast.Del): return tree raise SyntaxError("myield_from may only appear at the top level of a `@multishot` " "generator, as `myield_from[expr]` or `var = myield_from[expr]`") myield_from = namemacro(myield_from_function) # -------------------------------------------------------------------------------- # Expansion of `myield_from[expr]` / `var = myield_from[expr]` # # Architecture: # # - **Capture rest-of-outer at outer's top level**, before the iteration # begins, via `_rest = call_cc[get_cc()]`. The first pass through (when # `_rest` is a continuation) tail-calls the driver, passing `_rest`; the # second pass (when the driver has invoked `_rest(value)`) sees `_rest` as # the inner's `StopIteration` value and falls through to the post- # `myield_from` code in outer's body. # # - **The driver itself uses cut-the-tail** (`cc = identity` + return tuple) # to escape each `(captured_cc, inner_value)` to the user. Because the # driver is *tail-called* from outer (no nested trampoline started), and # the helper is *tail-called* via `call_cc[helper(...)]`, every step shares # the trampoline that `mi._k()` set up. The cut-the-tail escape therefore # reaches the user's `mi._k()` return — not a nested helper's frame. # # - **Resume after exhaustion via the captured rest-cc**: when inner raises # `StopIteration`, the driver does `return _rest_k(stopvalue)` (tail call), # resuming outer's body just after the rest-cc-capture point. # # - **`MultishotIterator` for inner**: convenient wrapper for `send`/`throw` # forwarding (we delegate to its protocol methods); it also lets # `gi_yieldfrom` surface the inner iterator via a stamp on the captured cc. # # Limitations of this v1: send/throw do reach the inner; `gi_yieldfrom` # tracks correctly while delegating; multi-shot fork during delegation is # inherited from the multi-shot semantics of the inner and the captured cc # stamping. See `doc/macros.md` for the user-facing description. def _build_myield_from_expansion(arg, target): """Build the AST list for a `myield_from` invocation. See the section comment above for the architecture. `arg` is the AST of the inner-multishot-call expression (e.g., `inner()`). `target` is the assignment target for `var = myield_from[...]`, or `None` for statement form (inner's `StopIteration` value discarded). """ inner_mi_name = gensym("_inner_mi") yieldf_name = gensym("_yieldf") drive_name = gensym("_drive") rest_name = gensym("_rest") with q as quoted: _INNER_MI_ = h[MultishotIterator](a[arg]) def _YIELDF_(_value, _inner_mi, *, cc): # Cut-the-tail: capture cc as `_k` (this is the at-call-cc # continuation, which is "rest of `_drive` after the call_cc"), # stamp it for `gi_yieldfrom`, then locally `cc = identity` so # the trampolined return delivers the (k, value) tuple straight # to whoever started the trampoline — i.e., the user's `mi._k()`. _k = cc _k._yieldfrom_inner = _inner_mi cc = h[identity] return (_k, _value) def _DRIVE_(_inner_mi, _rest_k, _value=None, _is_throw=False): try: if _is_throw: _x = _inner_mi.throw(_value) else: _x = _inner_mi.send(_value) except h[StopIteration] as _stopit: # Inner exhausted. Resume outer's body via the captured # rest-cc; outer continues at the post-`myield_from` code. return _rest_k(_stopit.value) _sent = h[call_cc][_YIELDF_(_x, _inner_mi)] if isinstance(_sent, BaseException) or h[safeissubclass](_sent, BaseException): return _DRIVE_(_inner_mi, _rest_k, _sent, True) return _DRIVE_(_inner_mi, _rest_k, _sent) # `_REST_ = call_cc[get_cc()]` is the multi-shot analog of Racket's # `(let/cc return ...)`: `_REST_` is bound to "rest of outer" as a # continuation. The first pass captures it; the driver later invokes # it with inner's `StopIteration` value to fall through to the rest. _REST_ = h[call_cc][h[get_cc]()] if h[iscontinuation](_REST_): return _DRIVE_(_INNER_MI_, _REST_) # control reaches here when inner exhausted; `_REST_` holds inner's # `StopIteration` value (or `None` if inner returned without a value). rename("_INNER_MI_", inner_mi_name, quoted) rename("_YIELDF_", yieldf_name, quoted) rename("_DRIVE_", drive_name, quoted) rename("_REST_", rest_name, quoted) if target is not None: with q as quoted_assign: a[target] = n[rest_name] quoted = quoted + quoted_assign return quoted # -------------------------------------------------------------------------------- # `@multishot` — decorator macro def multishot(tree, syntax, expander, **kw): """[syntax, decorator] Make a function into a multi-shot generator. Only meaningful inside a `with continuations:` block — required, not auto-wrapped. The expansion of `myield` produces `call_cc[get_cc()]`, which `with continuations` then turns into the actual continuation machinery; outside `with continuations`, that step fails with the standard `call_cc[]` SyntaxError. Multi-shot yield is spelled `myield`. The use site of `@multishot` must macro-import `myield` too, so that this macro knows which name you've bound it under. There are four variants:: Multi-shot yield Returns `k` expects Single-shot analog myield k no argument yield myield[expr] (k, value) no argument yield expr var = myield k one argument var = yield var = myield[expr] (k, value) one argument var = yield expr To resume, call `k`. In cases where `k` expects an argument, that argument is the value to send into `var`. Important differences from standard Python generators: - A multi-shot generator may be resumed from any `myield` arbitrarily many times, in any order. There is no concept of a single paused activation; each continuation is a function (technically a closure). When a multi-shot generator "myields", it returns just like a normal function, technically terminating its execution. But it gives you a continuation closure that you can call to resume execution just after that particular `myield`. The state lives in the closure cells of the continuation. The continuations are nested, so for a given activation, any locals in the already-executed part remain alive as long as at least one reference to a relevant continuation closure exists. "Nested" implies that re-invoking an earlier continuation branches execution into an independent timeline — but only for *locals*. Each resume gets a fresh activation record, so locals diverge; closure cells captured before the resume point are shared, so a mutation through one (a `nonlocal`, a mutable argument, or module-level state) is visible to every timeline reached from the same fork point. - `myield` is a *statement*, and it may only appear at the top level of a `@multishot` function definition (limitation of the underlying `call_cc[]`). Use inside lambdas, comprehensions, or nested `def`s is rejected at macro-expansion time. Usage:: with continuations: @multishot def f(): # Stop and return a continuation `k` that resumes just after this `myield`. myield # Stop and return the tuple `(k, 42)`. myield[42] # Stop and return a continuation `k`. Upon resuming `k`, # set the local `k` to the value sent in. k = myield # Stop and return the tuple `(k, 42)`. Upon resuming `k`, # set the local `k` to the value sent in. k = myield[42] # Instantiate the multi-shot generator (like calling a gfunc). # There is always an implicit bare `myield` at the beginning. k0 = f() # Start; run up to the explicit bare `myield`; receive new continuation. k1 = k0() # Continue to `myield[42]`; receive new continuation and the `42`. k2, x2 = k1() # Continue to `k = myield`; receive new continuation. k3 = k2() # Send `23` as the value of `k`; continue to `k = myield[42]`. k4, x4 = k3(23) # Send `17` as the value of `k`; continue to the end. # Reaching the end raises `StopIteration` (as with a regular generator). # `return value` inside `@multishot` raises `StopIteration(value)`. # Re-invoke an earlier continuation: k2, x2 = k1() For ergonomic generator-shaped consumption, wrap the initial continuation in a `MultishotIterator`. """ if syntax != "decorator": raise SyntaxError("multishot is a decorator macro only") # pragma: no cover if type(tree) is not ast.FunctionDef: raise SyntaxError("@multishot supports `def` only") # Detect the name(s) under which `myield` is macro-imported (handles as-imports). macro_bindings = extract_bindings(expander.bindings, myield_function) if not macro_bindings: raise SyntaxError("The use site of `@multishot` must macro-import `myield`, too.") names_of_myield = list(macro_bindings.keys()) # `myield_from` is optional; only present if the user macro-imported it. macro_bindings_from = extract_bindings(expander.bindings, myield_from_function) names_of_myield_from = set(macro_bindings_from.keys()) def is_myield_name(node): return type(node) is ast.Name and node.id in names_of_myield def is_myield_expr(node): return type(node) is ast.Subscript and is_myield_name(node.value) def is_myield_from_expr(node): # `myield_from[expr]` parses as a Subscript with `value` being the # `myield_from` name. Mirrors how `myield[expr]` is recognized. return (type(node) is ast.Subscript and type(node.value) is ast.Name and node.value.id in names_of_myield_from) def getslice(subscript_node): return subscript_node.slice class MultishotYieldTransformer(ASTTransformer): def transform(self, tree): if is_captured_value(tree): return tree if isnewscope(tree): return tree # `myield_from[expr]` / `var = myield_from[expr]` — handled before # `myield` shapes since both macros share the user-namespace symbol family. if names_of_myield_from: if (type(tree) is ast.Assign and len(tree.targets) == 1 and is_myield_from_expr(tree.value)): return _build_myield_from_expansion(getslice(tree.value), target=tree.targets[0]) if type(tree) is ast.Expr and is_myield_from_expr(tree.value): return _build_myield_from_expansion(getslice(tree.value), target=None) # `k = myield[value]` if type(tree) is ast.Assign and is_myield_expr(tree.value): if len(tree.targets) != 1: raise SyntaxError("expected exactly one assignment target in `k = myield[expr]`") var = tree.targets[0] value = getslice(tree.value) with q as quoted: a[var] = h[call_cc][h[get_cc]()] if h[iscontinuation](a[var]): return a[var], a[value] elif isinstance(a[var], BaseException) or h[safeissubclass](a[var], BaseException): raise a[var] return quoted # `k = myield` elif type(tree) is ast.Assign and is_myield_name(tree.value): if len(tree.targets) != 1: raise SyntaxError("expected exactly one assignment target in `k = myield`") var = tree.targets[0] with q as quoted: a[var] = h[call_cc][h[get_cc]()] if h[iscontinuation](a[var]): return a[var] elif isinstance(a[var], BaseException) or h[safeissubclass](a[var], BaseException): raise a[var] return quoted # `myield[value]` elif type(tree) is ast.Expr and is_myield_expr(tree.value): var = q[n[gensym("k")]] value = getslice(tree.value) with q as quoted: a[var] = h[call_cc][h[get_cc]()] if h[iscontinuation](a[var]): return h[partial](a[var], None), a[value] elif isinstance(a[var], BaseException) or h[safeissubclass](a[var], BaseException): raise a[var] return quoted # `myield` elif type(tree) is ast.Expr and is_myield_name(tree.value): var = q[n[gensym("k")]] with q as quoted: a[var] = h[call_cc][h[get_cc]()] if h[iscontinuation](a[var]): return h[partial](a[var], None) elif isinstance(a[var], BaseException) or h[safeissubclass](a[var], BaseException): raise a[var] return quoted return self.generic_visit(tree) class ReturnToRaiseStopIterationTransformer(ASTTransformer): def transform(self, tree): if is_captured_value(tree): return tree if isnewscope(tree): return tree if type(tree) is ast.Return: if tree.value is None: with q as quoted: raise h[StopIteration] return quoted with q as quoted: raise h[StopIteration](a[tree.value]) return quoted return self.generic_visit(tree) # Make the multishot generator raise `StopIteration` when it finishes via # any `return`. First make the implicit bare `return` explicit, then rewrite. # This must happen before transforming `myield`, to avoid breaking tail-calling # of the continuations. if type(tree.body[-1]) is not ast.Return: with q as quoted: return tree.body.extend(quoted) tree.body = ReturnToRaiseStopIterationTransformer().visit(tree.body) # Inject a bare `myield` resume point at the beginning of the function body. # When the multishot is initially called, the arguments are bound, and the # caller gets a continuation back; resuming that continuation actually starts # executing the function body. Mirrors a Python generator's first-`next` shape. tree.body.insert(0, ast.Expr(value=ast.Name(id=names_of_myield[0]))) tree.body = MultishotYieldTransformer().visit(tree.body) return tree # -------------------------------------------------------------------------------- # `MultishotIterator` — generator-protocol adapter def _continuation_code(k): """Extract the `__code__` of a continuation, unwrapping `partial` if present.""" func = k.func if isinstance(k, partial) else k return getattr(func, "__code__", None) def _step(k, mode, value=None): """Advance a continuation by one step. ``mode`` ∈ {"next", "send", "throw"}. ``"next"`` is treated as ``"send"`` with ``value=None`` — **matching the standard generator protocol**, where ``next(gen)`` is defined to be ``gen.send(None)``. For ``"send"``, if ``k`` is partial-wrapped (came from a bare ``myield`` or ``myield[expr]`` that doesn't bind a local), the sent value is dropped and the continuation advances normally — **also matching the standard generator protocol**, where ``gen.send(value)`` against a bare ``yield`` also discards the value silently. For raw-form continuations (from ``var = myield`` or ``var = myield[expr]``), the value is bound to ``var``. For ``"throw"``, the partial is unwrapped to inject the exception directly into the underlying continuation; the partial's pre-applied ``None`` would otherwise cause an arity mismatch. Returns whatever the continuation returns (typically a ``(next_k, value)`` tuple, or it raises if the continuation raises). """ if mode == "throw": underlying = k.func if isinstance(k, partial) else k return underlying(value) # mode in ("next", "send"); next ≡ send(None) if mode == "next": value = None return k() if isinstance(k, partial) else k(value) class MultishotIterator: """Adapt a `@multishot` generator to a subset of Python's generator protocol. Example:: with continuations: @multishot def g(): myield[1] myield[2] myield[3] mi = MultishotIterator(g()) assert [x for x in mi] == [1, 2, 3] Beyond the standard subset, `MultishotIterator` supports `copy.copy(mi)`, which forks the iterator at its current continuation. The fork shares the current continuation; subsequent advances get fresh activation records, so the forks' locals diverge. Closure cells captured before the fork point are shared, so a mutation through one (a `nonlocal`, a mutable argument, or module-level state) in one fork is visible to the others. Forks are independent timelines for *locals* only, not for state reached through closure cells. (Unlike standard generators, multi-shot generators support `copy.copy()`.) `copy.deepcopy(mi)` raises `TypeError` — the continuation closes over caller state we can't meaningfully deep-copy. Use `copy.copy(mi)` to fork. The current continuation is stored as `self.k` (read/write, type-checked). Overwriting `self.k` re-opens a closed iterator. Supported subset of the generator protocol: - `iter(mi)`, `next(mi)`, `mi.send(value)` - `mi.throw(exc)`, `mi.close()` - `mi.gi_code` — the `__code__` of the current continuation, or `None` when closed. **Use this as the liveness signal**, not `gi_frame`. - `mi.gi_frame` — **always `None`**. A multi-shot generator has no paused frame; state lives in the closure cells of the continuation. The standard-generator idiom `gen.gi_frame is None ↔ exhausted` does *not* apply here. - `mi.gi_running` — **always `False`**. Nothing is ever paused. - `mi.gi_yieldfrom` — currently always `None` (delegation via `myield_from` is not yet implemented). Not supported: - `yield from` across a standard generator and a multi-shot generator (semantic mismatch — standard generators have paused state, multi-shots don't; cannot be papered over). - Pickling — continuations are closures. - Async (`__aiter__`, `asend`, etc.). """ def __init__(self, k): self._k = None self._closed = False self.k = k # `self.k` — type-checked, fail-fast. @property def k(self): return self._k @k.setter def k(self, k): if not (iscontinuation(k) or (isinstance(k, partial) and iscontinuation(k.func))): raise TypeError( f"expected `k` to be a continuation or a partially applied continuation, got {k!r}" ) self._k = k self._closed = False # Generator-protocol introspection. @property def gi_frame(self): return None @property def gi_code(self): if self._closed: return None return _continuation_code(self._k) @property def gi_running(self): return False @property def gi_yieldfrom(self): if self._closed: return None # `_yieldfrom_inner` is stamped on the captured continuation by the # `myield_from` helper. The continuation is raw there (not partial- # wrapped), so a direct attribute read suffices — but check the # underlying function too for robustness in case a future code path # ever returns a partial-wrapped continuation from `myield_from`. underlying = self._k.func if isinstance(self._k, partial) else self._k return getattr(underlying, "_yieldfrom_inner", None) # Internal: drives one step via `_step`, updates `self._k` from the # returned `(next_k, value)`, and surfaces the value to the caller. # Used by `__next__`, `send`, and `throw`. def _advance(self, mode, value=None): assert mode in ("next", "send", "throw") if self._closed: raise StopIteration try: result = _step(self._k, mode, value) except StopIteration: self._closed = True raise if isinstance(result, tuple): self.k, x = result else: self.k, x = result, None return x # Generator API. def __iter__(self): return self def __next__(self): return self._advance("next") def send(self, value): return self._advance("send", value) def throw(self, exc): # Re-enters the current continuation, making it raise `exc`. If the # body catches and reaches another `myield`, the new continuation # becomes the current one and we return the next yielded value # (matching the standard generator protocol). If the exception isn't # caught, it propagates out of this call. return self._advance("throw", exc) def close(self): # https://docs.python.org/3/reference/expressions.html#generator.close # Bypass `_advance` here: close has different semantics (it injects # `GeneratorExit` and accepts `StopIteration` as a clean exit), and # `_advance` would short-circuit on the pre-set `_closed` flag. if self._closed: return self._closed = True try: _step(self._k, "throw", GeneratorExit) except GeneratorExit: return # body let the close exception propagate (expected) except StopIteration: return # body caught GeneratorExit and exited cleanly # Body caught `GeneratorExit` and `myield`ed another value — disallowed, # mirroring the standard generator protocol. raise RuntimeError("@multishot generator attempted to `myield` a value while it was being closed") # Forking — the multi-shot superpower exposed through the stdlib `copy` protocol. def __copy__(self): """Return a fork of this iterator at the current continuation. Both iterators share the current continuation. Each subsequent advance gets a fresh activation record, so the forks' locals diverge. Closure cells captured before the fork point are shared, so a mutation through one (a `nonlocal`, a mutable argument, or module-level state) in one fork is visible to the others. Forks are independent timelines for *locals* only, not for state reached through closure cells. This is the multi-shot semantics, not a quirk of `copy.copy()`; the same applies to plain re-invocation of an earlier continuation. If this iterator is closed, the fork is also closed (the underlying continuation is preserved, so the fork can be re-opened by assigning to `.k`). """ forked = MultishotIterator(self._k) forked._closed = self._closed return forked def __deepcopy__(self, memo): raise TypeError( "multi-shot iterators cannot be deep-copied; use copy.copy() to fork" ) def __del__(self): # Mirror generator GC semantics. Mostly cosmetic for multishots # (no paused frame to clean up), but politely closes the iterator. try: self.close() except Exception: pass