"""
Subscript-based macro expansion implemented with Pyccolo, to accompany PipelineTracer.
"""
from __future__ import annotations
import ast
import builtins
import re
import textwrap
from contextlib import contextmanager
from functools import reduce
from types import FrameType
from typing import Any, Callable, Generator, cast
import pyccolo as pyc
from pyccolo import fast
from pyccolo.stmt_mapper import StatementMapper
from pyccolo.trace_events import TraceEvent
from typing_extensions import Literal
import pipescript.api.static_macros
from pipescript.analysis.dynamic_macros import DynamicMacro
from pipescript.analysis.placeholders import SingletonArgCounterMixin
from pipescript.api.static_macros import (
_ntimes_counters,
_once_cache,
context,
do,
expect,
fork,
future,
memoize,
ntimes,
once,
otherwise,
parallel,
read,
repeat,
unless,
until,
when,
write,
)
from pipescript.tracers.pipeline_tracer import PipelineTracer
from pipescript.utils import get_user_ns
class ArgReplacer(ast.NodeVisitor, SingletonArgCounterMixin):
def __init__(self) -> None:
super().__init__()
self.placeholder_names: list[str] = []
self.arg_node_id_to_placeholder_name: dict[int, str] = {}
self._macro_visit_context: bool = False
@contextmanager
def macro_visit_context(self, override: bool = True) -> Generator[None, None, None]:
old = self._macro_visit_context
try:
self._macro_visit_context = override
yield
finally:
self._macro_visit_context = old
@property
def placeholder_nodes(self) -> set[int]:
if self._macro_visit_context:
return PipelineTracer.augmented_node_ids_by_spec[
PipelineTracer.macro_arg_placeholder_spec
]
else:
return PipelineTracer.augmented_node_ids_by_spec[
PipelineTracer.arg_placeholder_spec
]
def visit_Subscript(self, node: ast.Subscript) -> None:
if (
not self._macro_visit_context
and isinstance(node.value, ast.Name)
and node.value.id in MacroTracer.static_macros
):
# defer visiting nested quick lambdas
return
self.generic_visit(node)
def visit_Name(self, node: ast.Name) -> None:
if id(node) not in self.placeholder_nodes:
return
self.placeholder_nodes.discard(id(node))
assert node.id.startswith("_")
if node.id == "_":
node.id = f"_{self.arg_ctr}"
self.arg_ctr += 1
else:
if not node.id[1].isdigit():
node.id = node.id[1:]
else:
# ensure we prevent collisions with auto placeholders like _1, _2, etc
node.id = "_" + node.id
if node.id not in self.placeholder_names:
self.placeholder_names.append(node.id)
self.arg_node_id_to_placeholder_name[id(node)] = node.id
def visit_pipeline(self, node: ast.BinOp) -> None:
num_left_traversals = PipelineTracer.search_left_descendant_placeholder(node)
if num_left_traversals < 0:
return
left_arg: ast.expr = node
for _ in range(num_left_traversals):
left_arg = left_arg.left # type: ignore[attr-defined]
self.visit(left_arg)
def visit_BinOp(self, node: ast.BinOp) -> None:
if (
not self._macro_visit_context
and isinstance(node.op, ast.BitOr)
and PipelineTracer.get_augmentations(id(node))
):
self.visit_pipeline(node)
else:
self.generic_visit(node)
def __call__(self, node: ast.AST) -> None:
self.arg_node_id_to_placeholder_name.clear()
self.placeholder_names.clear()
self.visit(node)
def get_placeholder_names(self, node: ast.AST) -> list[str]:
self(node)
return self.placeholder_names
#: Name of the sentinel builtin a statement-block marker calls (see
#: :func:`block_marker_id` / :func:`_block_marker_sentinel`).
BLOCK_MARKER_FUNC = "__pyc_block__"
def _block_marker_sentinel(_n: int) -> None:
"""Resolve a statement-block marker (``__pyc_block__(N)``) to ``None``.
The marker is emitted as a call to this *defined* builtin rather than a bare
(undefined) ``__pyc_block_N__`` name. ipyflow's ``before_subscript_slice``
machinery evaluates the original slice expression *before* the macro handler
can substitute it, so an undefined marker name leaks as a ``NameError``
before ``handle_macro`` ever runs. Calling a defined sentinel evaluates
harmlessly to ``None``; the handler (which keys off the AST, not the value)
then fires and replaces the slice with the compiled block, exactly as it
does for an ordinary ``map[int]``-style slice. (Base pyccolo substitutes
lazily and never evaluated the marker, which is why this only bit ipyflow.)
"""
return None
# Register the sentinel as a builtin so the marker call resolves anywhere the
# slice is evaluated (cell globals, compiled-block globals) and survives the
# per-cell tracer ``reset()`` that strips macro-name builtins.
setattr(builtins, BLOCK_MARKER_FUNC, _block_marker_sentinel)
def block_marker_id(node: ast.AST) -> int | None:
"""If ``node`` is a statement-block macro slice (``macro[__pyc_block__(N)]``),
return ``N``; otherwise ``None``."""
if not isinstance(node, ast.Subscript):
return None
slc = node.slice
if (
isinstance(slc, ast.Call)
and isinstance(slc.func, ast.Name)
and slc.func.id == BLOCK_MARKER_FUNC
and len(slc.args) == 1
and isinstance(slc.args[0], ast.Constant)
and isinstance(slc.args[0].value, int)
):
return slc.args[0].value
return None
class _BlockPlaceholderReplacer(ast.NodeVisitor):
"""Identify ``$`` placeholders in a (already ``$``->``_`` transformed)
statement block by name convention and rename them to concrete parameter
names, mirroring :class:`ArgReplacer` but for statement bodies whose nodes
were not position-marked.
Unlike the expression form -- where each bare ``$`` is a *fresh* positional
argument -- a statement block treats every bare ``$`` as the single piped
input (parameter ``_0``), so it can be referenced repeatedly. Use named
placeholders (``$name``) for additional distinct parameters."""
def __init__(self) -> None:
self.auto: list[str] = []
self.named: list[str] = []
@property
def params(self) -> list[str]:
return self.auto + self.named
def visit_Subscript(self, node: ast.Subscript) -> None:
# defer nested macros -- their placeholders belong to them
if isinstance(node.value, ast.Name) and node.value.id in (
MacroTracer.static_macros | MacroTracer.dynamic_macros
):
return
self.generic_visit(node)
def visit_Name(self, node: ast.Name) -> None:
name = node.id
if name == "_":
node.id = "_0"
if "_0" not in self.auto:
self.auto.append("_0")
elif len(name) > 1 and name[0] == "_" and name[1].isalpha():
pname = name[1:]
node.id = pname
if pname not in self.named:
self.named.append(pname)
def normalize_block_src(src: str) -> str:
"""Normalize a captured block body to consistent indentation.
The body's first statement is typically *inline* with the opening ``{`` (so
it carries the column of the ``{``), while later statements carry the source
indentation -- which an enclosing block's dedent may have shifted relative to
the first line. ``textwrap.dedent`` can't fix that (the minimum indent may be
on a later line). Instead we dedent every line by the *first* line's indent,
clamping at zero, which puts the first statement at column 0 and preserves
the relative indentation of nested suites beneath it."""
lines = src.split("\n")
while lines and not lines[0].strip():
lines.pop(0)
while lines and not lines[-1].strip():
lines.pop()
if not lines:
return ""
base = len(lines[0]) - len(lines[0].lstrip())
out = []
for ln in lines:
if not ln.strip():
out.append("")
continue
indent = len(ln) - len(ln.lstrip())
out.append(" " * max(0, indent - base) + ln.lstrip())
return "\n".join(out)
# Forward-pipe operators (``|>`` and friends, with optional ``*`` / ``**``
# unpack prefixes). Their right-hand *stage* expression has its own ``$``
# placeholders (the piped value), which belong to PipelineTracer -- not to the
# block's collapse-``$`` semantics. Matched longest-first so e.g. ``**|>`` wins
# over ``|>``.
_FORWARD_PIPE_RE = re.compile(r"(?:\*\*|\*)?(?:\|>>|\|>|\$>|\?>|\.>)")
# Tokens that end a pipe stage, so a ``$`` after them is the block's again. We
# scope conservatively: commas, comparisons (lower precedence than ``|``), the
# ternary/boolean keywords, and assignment/colon. (The ``<``/``>`` that form
# pipe operators are consumed as pipe spans below, so a bare ``<``/``>`` here is
# a real comparison.)
_STAGE_TERMINATOR_OPS = frozenset(
{",", "==", "!=", "<", ">", "<=", ">=", "=", ":", ";"}
)
_STAGE_TERMINATOR_KW = frozenset({"and", "or", "not", "in", "is", "if", "else", "for"})
def split_block_placeholders(src: str) -> tuple[str, list[str]]:
"""Replace the *block's own* ``$`` placeholders with parameter names while
leaving placeholders that belong to nested macros (``macro[...]`` /
``macro{...}``) or to nested *pipelines* (the stage after ``|>``) untouched,
so they are processed by those handlers later.
Returns ``(new_src, params)``. Bare ``$`` collapses to the single piped
input ``_0``; ``$name`` becomes a distinct parameter ``name``. A ``$`` is
"nested" when it sits inside a ``[``/``{`` that immediately follows a macro
name, or in a forward-pipe stage at the current bracket depth (e.g. the
second ``$`` in ``$ |> $ + 1`` is the pipe's argument, not the block's)."""
import io
import keyword as _kw
import tokenize as _tok
macro_names = (
set(MacroTracer.static_macros)
| set(MacroTracer.dynamic_macros)
| set(MacroTracer.dynamic_method_macros)
)
try:
toks = list(_tok.generate_tokens(io.StringIO(src).readline))
except (_tok.TokenError, IndentationError, SyntaxError):
return src, []
line_starts = [0]
for i, ch in enumerate(src):
if ch == "\n":
line_starts.append(i + 1)
def _off(pos: tuple[int, int]) -> int:
row = pos[0] - 1
if row >= len(line_starts):
return len(src)
return line_starts[row] + pos[1]
pipe_spans = [m.span() for m in _FORWARD_PIPE_RE.finditer(src)]
def _in_pipe_span(start: int, end: int) -> bool:
return any(s < end and start < e for s, e in pipe_spans)
repls: list[tuple[int, int, str]] = []
auto: list[str] = []
named: list[str] = []
macro_depth = 0
opener_is_macro: list[bool] = []
# one "pipe stage seen" flag per open-bracket depth (index 0 == top level),
# so a pipe inside ``(...)`` doesn't leak out to siblings.
pipe_seen: list[bool] = [False]
prev = None
for i, t in enumerate(toks):
if t.type in (
_tok.NL,
_tok.NEWLINE,
_tok.INDENT,
_tok.DEDENT,
_tok.COMMENT,
_tok.ENCODING,
_tok.ENDMARKER,
):
if t.type in (_tok.NL, _tok.NEWLINE):
pipe_seen[-1] = False
continue
start, end = _off(t.start), _off(t.end)
if _in_pipe_span(start, end):
# part of a forward-pipe operator: opens a stage at this depth
pipe_seen[-1] = True
prev = t
continue
if t.type == _tok.OP and t.string in ("(", "[", "{"):
if t.string in ("[", "{"):
is_macro = (
prev is not None
and prev.type == _tok.NAME
and prev.string in macro_names
and prev.end == t.start
)
opener_is_macro.append(is_macro)
macro_depth += int(is_macro)
pipe_seen.append(False)
elif t.type == _tok.OP and t.string in (")", "]", "}"):
if t.string in ("]", "}") and opener_is_macro:
macro_depth -= int(opener_is_macro.pop())
if len(pipe_seen) > 1:
pipe_seen.pop()
elif (t.type == _tok.OP and t.string in _STAGE_TERMINATOR_OPS) or (
t.type == _tok.NAME and t.string in _STAGE_TERMINATOR_KW
):
pipe_seen[-1] = False
elif t.string == "$" and macro_depth == 0 and not pipe_seen[-1]:
nxt = toks[i + 1] if i + 1 < len(toks) else None
if (
nxt is not None
and nxt.type == _tok.NAME
and not _kw.iskeyword(nxt.string)
and t.end == nxt.start
):
pname = nxt.string
repls.append((start, _off(nxt.end), pname))
if pname not in named:
named.append(pname)
else:
repls.append((start, end, "_0"))
if "_0" not in auto:
auto.append("_0")
prev = t
for rstart, rend, text in sorted(repls, reverse=True):
src = src[:rstart] + text + src[rend:]
return src, auto + named
def _block_assigned_names(body: list[ast.stmt]) -> list[str]:
"""Top-level names a namespace block binds, in source order (first occurrence
sets position). ``_``-prefixed names are treated as block-local temporaries
and excluded, so e.g. ``_scale = 0.1; w = _scale * ...`` harvests only ``w``."""
names: list[str] = []
seen: set[str] = set()
def add(name: str) -> None:
if name.startswith("_") or name in seen:
return
seen.add(name)
names.append(name)
for stmt in body:
if isinstance(stmt, ast.Assign):
for tgt in stmt.targets:
if isinstance(tgt, ast.Name):
add(tgt.id)
elif isinstance(tgt, (ast.Tuple, ast.List)):
for elt in tgt.elts:
if isinstance(elt, ast.Name):
add(elt.id)
elif isinstance(stmt, ast.AnnAssign) and isinstance(stmt.target, ast.Name):
if stmt.value is not None:
add(stmt.target.id)
elif isinstance(stmt, ast.AugAssign) and isinstance(stmt.target, ast.Name):
add(stmt.target.id)
return names
[docs]
def register_namespace_macro(
name: str,
builder: Callable[[dict], Any],
call_form: Callable[..., Any] | None = None,
) -> None:
"""Register ``name{ a = ...; b = ... }`` as a namespace-block macro.
The brace block is run and its top-level assignments harvested into a dict,
which is passed to ``builder``; the macro evaluates to ``builder(namespace)``.
``call_form``, if given, is exposed as the bare name ``name`` (a builtin in a
pipescript session) so the ordinary call form ``name(a=..., b=...)`` works too;
pass the function that ``builder`` delegates to. Registration must happen
before the cell/code using the macro is processed.
"""
MacroTracer.namespace_block_macros[name] = builder
MacroTracer.static_macros[name] = call_form
# The macro name is still evaluated as an ordinary load before the
# subscript-load handler swaps in the macro machinery, so it must resolve to
# *something*; bind it as a builtin (the call form, or None as a placeholder)
# exactly as ``MacroTracer.__init__`` does for the built-in macros.
setattr(builtins, name, call_form)
def is_static_macro(node: ast.AST) -> bool:
return (
isinstance(node, ast.Subscript)
and isinstance(node.value, ast.Name)
and node.value.id in MacroTracer.static_macros
)
def is_dynamic_macro(node: ast.AST) -> bool:
return (
isinstance(node, ast.Subscript)
and isinstance(node.value, ast.Name)
and node.value.id in MacroTracer.dynamic_macros
)
def is_dynamic_method_macro_attribute(node: ast.AST) -> bool:
return (
isinstance(node, ast.Attribute)
and node.attr in MacroTracer.dynamic_method_macros
)
def is_dynamic_method_macro(node: ast.AST) -> bool:
return isinstance(node, ast.Subscript) and is_dynamic_method_macro_attribute(
node.value
)
class MacroTracer(pyc.BaseTracer):
allow_reentrant_events = True
global_guards_enabled = False
multiple_threads_allowed = True
# macros ending in 'f' allow us to write stuff like reduce[+]([1, 2, 3])
static_macros = {
context.__name__: context,
do.__name__: do,
expect.__name__: expect,
"f": None,
fork.__name__: fork,
future.__name__: future,
filter.__name__: filter,
f"i{filter.__name__}": filter,
f"{filter.__name__}f": filter,
f"i{filter.__name__}f": filter,
"macro": None,
map.__name__: map,
f"i{map.__name__}": map,
f"{map.__name__}f": map,
f"i{map.__name__}f": map,
memoize.__name__: memoize,
"method": None,
ntimes.__name__: ntimes,
once.__name__: once,
otherwise.__name__: otherwise,
parallel.__name__: parallel,
read.__name__: read,
reduce.__name__: reduce,
f"{reduce.__name__}f": reduce,
repeat.__name__: repeat,
unless.__name__: unless,
until.__name__: until,
when.__name__: when,
write.__name__: write,
}
dynamic_macros: dict[str, DynamicMacro] = {}
dynamic_method_macros: dict[str, DynamicMacro] = {}
#: Namespace-block macros: ``name`` -> builder. ``name{ a = ...; b = ... }``
#: runs the block, harvests its top-level assignments into a dict, and returns
#: ``builder(that_dict)``. Registered via :func:`register_namespace_macro`
#: (e.g. pyccolo's autodiff example wires ``params`` to its ``Param`` builder).
namespace_block_macros: dict[str, Callable[[dict], Any]] = {}
builtin_dynamic_macro_definitions: dict[str, str] = {
"foreach": "method[$$ |> map[do[$$]] |> list]",
}
assert set(pipescript.api.static_macros.__all__) <= set(static_macros.keys())
_not_found = object()
def __init__(self, *args, **kwargs) -> None:
super().__init__(*args, **kwargs)
self.arg_replacer = ArgReplacer()
# (id(node), id(frame), evt) -> (frame, compiled thunk). The frame is kept
# only so a lookup can confirm identity: id(frame) is recycled once a frame
# is collected, and returning a *baked* macro thunk from a prior evaluation
# whose frame id was reused would skip re-execution. Retention is bounded
# by evicting entries whose frame has left the live call stack (see
# ``_lambda_cache_lookup``) -- without that, holding a frame would pin its
# locals, e.g. a whole training-step autograd graph, and leak per step.
self.lambda_cache: dict[tuple[int, int, TraceEvent], tuple[FrameType, Any]] = {}
self._overridden_builtins: list[str] = []
# An exception raised by user code inside a namespace block, tagged so
# ``should_propagate_handler_exception`` lets it surface verbatim rather
# than being swallowed into a ``None`` handler result (see that method).
self.exc_to_propagate: Exception | None = None
user_ns = get_user_ns()
for macro_name, macro in (
self.static_macros | self.dynamic_macros | self.dynamic_method_macros
).items():
if hasattr(builtins, macro_name):
continue
setattr(builtins, macro_name, macro)
self._overridden_builtins.append(macro_name)
if user_ns is not None:
user_ns.setdefault(macro_name, macro)
def reset(self) -> None:
for macro_name in self._overridden_builtins:
if hasattr(builtins, macro_name):
delattr(builtins, macro_name)
self._overridden_builtins.clear()
self.lambda_cache.clear()
super().reset()
def _lambda_cache_lookup(
self, node: ast.AST, frame: FrameType, evt: TraceEvent
) -> Any:
"""Return the cached thunk for ``(node, frame, evt)`` or ``_not_found``.
First evict every entry whose frame is no longer on ``frame``'s live call
stack: that keeps the cache from pinning a completed evaluation's frame
(and its locals) while still serving the *re-entrant* lookups within a
single evaluation, which share one live frame. The identity check guards
against ``id(frame)`` reuse handing back a prior frame's baked thunk."""
if self.lambda_cache:
live: set[FrameType] = set()
f: FrameType | None = frame
while f is not None:
live.add(f)
f = f.f_back
stale = [
key
for key, (cached_frame, _) in self.lambda_cache.items()
if cached_frame not in live
]
for key in stale:
del self.lambda_cache[key]
entry = self.lambda_cache.get((id(node), id(frame), evt))
if entry is not None and entry[0] is frame:
return entry[1]
return self._not_found
def _lambda_cache_store(
self, node: ast.AST, frame: FrameType, evt: TraceEvent, thunk: Any
) -> Any:
self.lambda_cache[(id(node), id(frame), evt)] = (frame, thunk)
return thunk
def should_propagate_handler_exception(
self, _evt: TraceEvent, exc: Exception
) -> bool:
# pyccolo defensively swallows handler exceptions (so an instrumentation
# bug can't break user code). A namespace block, however, runs user code:
# let an exception we tagged in ``_handle_namespace_block`` propagate so
# the user sees their real traceback, not a None context value.
if exc is self.exc_to_propagate:
self.exc_to_propagate = None
return True
return False
class _IdentityAttributeSubscript:
def __getitem__(self, item):
return item
def __getattr__(self, item):
return self
_identity_attribute_subscript = _IdentityAttributeSubscript()
@pyc.before_attribute_load(when=is_dynamic_method_macro_attribute, reentrant=True)
def ignore_dynamic_macro_attribute_load(self, *_, **__):
return self._identity_attribute_subscript
@pyc.before_subscript_load(
when=lambda node: is_dynamic_macro(node) or is_dynamic_method_macro(node),
reentrant=True,
)
def ignore_dynamic_macro_slice(self, *_, **__):
return self._identity_attribute_subscript
@pyc.before_subscript_slice(
when=lambda node: is_dynamic_macro(node) or is_dynamic_method_macro(node),
reentrant=True,
)
def perform_dynamic_macro_substitution(
self,
_ret,
node: ast.Subscript,
frame: FrameType,
evt: TraceEvent,
*_,
**__,
):
cached_lambda = self._lambda_cache_lookup(node, frame, evt)
if cached_lambda is not self._not_found:
return cached_lambda
__hide_pyccolo_frame__ = True
if isinstance(node.value, ast.Name):
macro_instance = self.dynamic_macros[node.value.id]
macro_body = node.slice
elif isinstance(node.value, ast.Attribute):
macro_instance = self.dynamic_method_macros[node.value.attr]
with fast.location_of(node.slice):
macro_body = fast.Tuple(elts=[node.value.value, node.slice])
else:
raise ValueError(
"Impossible node type for dynamic macro substitution: %s"
% type(node.value)
)
expanded_macro_expr = macro_instance.expand(macro_body)
if isinstance(expanded_macro_expr, ast.expr):
lambda_body = expanded_macro_expr
orig_ctr = self.arg_replacer.arg_ctr
self.arg_replacer.arg_ctr = orig_ctr
expanded_macro_expr, *_extra = self.arg_replacer.create_placeholder_lambda(
[],
self.arg_replacer.arg_ctr,
expanded_macro_expr,
frame,
)
expanded_macro_expr.body = lambda_body
callable_func = pyc.eval(
expanded_macro_expr, frame.f_globals, frame.f_locals
)
try:
evaluated_lambda = callable_func()
except Exception as e:
exc = e
def raises_error():
raise exc
return lambda: __hide_pyccolo_frame__ and raises_error()
ret = lambda: __hide_pyccolo_frame__ and evaluated_lambda # noqa: E731
else:
evaluated_lambda = expanded_macro_expr
ret = lambda: __hide_pyccolo_frame__ and evaluated_lambda # noqa: E731
return self._lambda_cache_store(node, frame, evt, ret)
@pyc.before_subscript_load(when=is_static_macro, reentrant=True)
def load_macro_result(self, *_, **__):
return self._identity_attribute_subscript
def _transform_ast_lambda_for_macro(
self,
ast_lambda: ast.expr,
func: str,
extra_defaults: set[str],
) -> ast.Lambda:
arg = f"_{self.arg_replacer.arg_ctr}"
self.arg_replacer.arg_ctr += 1
starred_arg = f"_{self.arg_replacer.arg_ctr}"
self.arg_replacer.arg_ctr += 1
inner_func = func
if func == "ifilter":
inner_func = "filter"
elif func == "imap":
inner_func = "map"
lambda_body_str = f"{inner_func}(None, {arg}, *{starred_arg})"
functor_lambda_body = cast(
ast.Call,
cast(
ast.Expr,
fast.parse(lambda_body_str).body[0],
).value,
)
functor_lambda_body.args[0] = ast_lambda
if func in ("filter", "map"):
id_arg = f"_{self.arg_replacer.arg_ctr}"
self.arg_replacer.arg_ctr += 1
lambda_body_str = f"(type({arg}) if type({arg}) in (frozenset, list, set, tuple) else lambda {id_arg}: {id_arg})(None)"
functor_lambda_outer_body = cast(
ast.Call,
cast(
ast.Expr,
fast.parse(lambda_body_str).body[0],
).value,
)
functor_lambda_outer_body.args[0] = functor_lambda_body
functor_lambda_body = functor_lambda_outer_body
functor_lambda = cast(
ast.Lambda,
cast(
ast.Expr,
fast.parse(
f"lambda {arg}, *{starred_arg}, {', '.join(f'{name}={name}' for name in extra_defaults)}: None"
).body[0],
).value,
)
functor_lambda.body = functor_lambda_body
return functor_lambda
def _handle_macro_impl(
self,
orig_lambda_body: ast.expr,
frame: FrameType,
func: str,
allow_call_node: bool = True,
) -> ast.expr:
__hide_pyccolo_frame__ = True # noqa: F841
orig_ctr = self.arg_replacer.arg_ctr
lambda_body = StatementMapper.bookkeeping_propagating_copy(orig_lambda_body)
placeholder_names = self.arg_replacer.get_placeholder_names(lambda_body)
needs_call_node = (
self.arg_replacer.arg_ctr == orig_ctr and len(placeholder_names) == 0
)
(
ast_lambda,
extra_defaults,
modified_lambda_body,
) = SingletonArgCounterMixin.create_placeholder_lambda(
placeholder_names,
orig_ctr,
lambda_body,
frame,
created_starred_arg=needs_call_node,
)
lambda_body = modified_lambda_body or lambda_body
if needs_call_node and allow_call_node:
with fast.location_of(lambda_body):
load = ast.Load()
lambda_body = fast.Call(
func=lambda_body,
args=[
fast.Starred(
fast.Name(f"_{self.arg_replacer.arg_ctr - 1}", ctx=load),
ctx=load,
)
],
keywords=[],
)
ast_lambda.body = lambda_body
if func in self.static_macros and func not in (
"f",
memoize.__name__,
otherwise.__name__,
):
with fast.location_of(ast_lambda):
ast_lambda = self._transform_ast_lambda_for_macro(
ast_lambda, func, extra_defaults
)
ret_expr: ast.expr = ast_lambda
if func in (memoize.__name__, otherwise.__name__):
with fast.location_of(ast_lambda):
ret_expr = fast.Call(
func=fast.Name(func, ctx=ast.Load()),
args=[ast_lambda],
keywords=[],
)
return ret_expr
@fast.location_of_arg
def _handle_read_write_macro(
self, func: Literal["read", "write"], macro_body: ast.expr
) -> ast.Lambda:
if isinstance(macro_body, ast.Name):
macro_body = fast.Str(macro_body.id)
arg = f"_{self.arg_replacer.arg_ctr}"
self.arg_replacer.arg_ctr += 1
lam: ast.Lambda = cast(
ast.Lambda,
cast(
ast.Expr, fast.parse(f"lambda {arg}: {func}(None, {arg})").body[0]
).value,
)
cast(ast.Call, lam.body).args[0] = macro_body
return lam
@fast.location_of_arg
def _handle_ntimes_macro(
self, frame: FrameType, macro_body: ast.expr
) -> ast.Lambda:
callpoint_id = id(macro_body)
if callpoint_id not in _ntimes_counters:
if isinstance(macro_body, ast.Constant) and isinstance(
macro_body.value, int
):
ctr = macro_body.value
else:
ctr = pyc.eval(macro_body, frame.f_globals, frame.f_locals)
_ntimes_counters[callpoint_id] = ctr
arg = f"_{self.arg_replacer.arg_ctr}"
self.arg_replacer.arg_ctr += 1
lam: ast.Lambda = cast(
ast.Lambda,
cast(
ast.Expr,
fast.parse(f"lambda {arg}: {ntimes.__name__}({arg}, None)").body[0],
).value,
)
cast(ast.Call, lam.body).args[1] = fast.Constant(value=callpoint_id)
return lam
@fast.location_of_arg
def _handle_once_macro(self, frame: FrameType, macro_body: ast.expr) -> ast.expr:
callpoint_id = id(macro_body)
if callpoint_id not in _once_cache:
_once_cache[callpoint_id] = pyc.eval(
macro_body, frame.f_globals, frame.f_locals
)
expr = cast(
ast.Expr,
fast.parse(f"{once.__name__}(None)").body[0],
).value
cast(ast.Call, expr).args[0] = fast.Constant(value=callpoint_id)
return expr
def _block_free_vars(
self, body: list[ast.stmt], params: set[str], frame: FrameType
) -> set[str]:
"""Names the block reads but never assigns (and that aren't params,
builtins, or globals); resolved against the defining frame's f_back
chain. Computed from the placeholder-substituted source before
instrumentation so pyccolo's own instrumentation names don't leak in."""
loads: set[str] = set()
stores: set[str] = set()
class _NameCollector(ast.NodeVisitor):
def visit_Name(self, n: ast.Name) -> None:
(stores if isinstance(n.ctx, ast.Store) else loads).add(n.id)
collector = _NameCollector()
for stmt in body:
collector.visit(stmt)
return {
n
for n in loads - stores - params
if not hasattr(builtins, n) and not n.startswith("__pyc_")
# names already visible as globals/builtins resolve directly;
# `_dynamic_lookup` only walks enclosing *locals*.
and n not in frame.f_globals
}
@staticmethod
def _name_in_enclosing_locals(frame: FrameType, name: str) -> bool:
"""Whether ``name`` is bound as a *function* local in some enclosing
(non-machinery) frame. Machinery frames mark themselves with
``__hide_pyccolo_frame__`` and are skipped, matching the runtime resolver
in ``pipescript.api.utils._resolve_enclosing_frame``; module frames
(``f_locals is f_globals``) are excluded so only genuine locals qualify."""
fr: FrameType | None = frame
while fr is not None:
if (
"__hide_pyccolo_frame__" not in fr.f_locals
and fr.f_locals is not fr.f_globals
and name in fr.f_locals
):
return True
fr = fr.f_back
return False
def _block_enclosing_assigned(
self, body: list[ast.stmt], params: set[str], frame: FrameType
) -> set[str]:
"""Top-level names the block assigns that already exist as an enclosing
function local -- these get read-back-write-back semantics so the block
can rebind them (e.g. ``total = total + $``). Only *statement-level*
assignment targets are considered; names bound in nested scopes
(comprehensions, nested ``def``/``lambda``) stay local to those scopes."""
assigned: set[str] = set()
def add_target(tgt: ast.expr) -> None:
if isinstance(tgt, ast.Name):
assigned.add(tgt.id)
elif isinstance(tgt, (ast.Tuple, ast.List)):
for elt in tgt.elts:
add_target(elt)
elif isinstance(tgt, ast.Starred):
add_target(tgt.value)
for stmt in body:
if isinstance(stmt, ast.Assign):
for tgt in stmt.targets:
add_target(tgt)
elif isinstance(stmt, (ast.AugAssign, ast.AnnAssign)):
add_target(stmt.target)
return {
n
for n in assigned - params
if not hasattr(builtins, n)
and not n.startswith("__pyc_")
and n not in frame.f_globals
and self._name_in_enclosing_locals(frame, n)
}
def _compile_block_function(
self, block_src: str, frame: FrameType, namespace: bool = False
):
"""Compile a stashed statement block into a function: the block's own
``$`` placeholders become parameters, the trailing expression becomes the
return value, free variables resolve against the defining frame, and
*nested* pipescript syntax (``|>``, ``f[...]``, ``f{...}``) is parsed,
marked, and instrumented so it dispatches at runtime.
With ``namespace=True`` the function instead returns a ``dict`` of the
block's top-level assignments -- ``name = expr`` lines, excluding
``_``-prefixed temporaries -- so a brace macro can harvest the block as a
namespace (e.g. ``params{ w = ...; b = ... }`` -> ``{"w": ..., "b": ...}``).
The body is compiled via :meth:`~pyccolo.BaseTracer.parse_fragment`,
which instruments the fragment from inside this active macro expansion
without re-entering / corrupting it."""
# Hidden so a co-tracer (e.g. ipyflow) that walks the frame stack to map
# a sandbox call back to a notebook position skips this frame instead of
# misattributing its own line number to the executing cell.
__hide_pyccolo_frame__ = True # noqa: F841
from pipescript.analysis.placeholders import FreeVarTransformer
from pipescript.api.utils import _dynamic_lookup, _dynamic_store
block_src = normalize_block_src(block_src)
# Substitute the block's own placeholders for parameters, leaving nested
# macros' placeholders (still `$`) for pyccolo to mark/instrument.
param_src, params = split_block_placeholders(block_src)
param_src = normalize_block_src(param_src)
# Free-var analysis on a clean (uninstrumented) parse of the same source.
clean_body = ast.parse(textwrap.dedent(self.transform(param_src)).strip("\n"))
freevars = self._block_free_vars(clean_body.body, set(params), frame)
# Names the block *assigns* that already exist as an enclosing function
# local: rebinding one should write back to that local, mirroring inline
# code. A namespace block instead harvests its assignments into a dict, so
# it keeps its assignments block-local.
enclosing_assigned = (
set()
if namespace
else self._block_enclosing_assigned(clean_body.body, set(params), frame)
)
# A name resolved for write-back must not also be rewritten as a read-only
# free var (it is a real block local, seeded and flushed below).
freevars = freevars - enclosing_assigned
# Parse the body as a *module* (at column 0) and wrap it in a FunctionDef
# at the AST level. Wrapping in `def ...:` *textually* (with
# textwrap.indent) corrupts a nested block whose first statement is inline
# with `{` -- its lines end up inconsistently indented and fail to parse.
#
# Scope the instrumentation to the substituting tracers (those with
# ``global_guards_enabled = False`` -- i.e. pipescript's own, which the
# block's nested macros need), excluding pure observers like ipyflow's
# dataflow tracer. The block is synthetic sandbox code an observer cannot
# map back to a notebook statement; weaving its statement events in makes
# the observer raise on the block's unknown nodes.
from pyccolo.tracer import _TRACER_STACK
block_tracers = [
t for t in _TRACER_STACK if not t.global_guards_enabled
] or None
body_module = cast(
ast.Module, self.parse_fragment(param_src, tracers=block_tracers)
)
body = list(body_module.body)
# instrumentation may wrap the body in a try/except; the trailing
# expression we want to return lives at the innermost level.
target_body = body
while len(target_body) == 1 and isinstance(target_body[0], ast.Try):
target_body = target_body[0].body
if namespace:
# Return a dict of the block's top-level assignments (collected from
# the clean, pre-instrumentation parse so pyccolo temporaries do not
# leak in). The Name loads reference block-local bindings, so they are
# not free vars and the FreeVarTransformer below leaves them alone.
names = _block_assigned_names(clean_body.body)
target_body.append(
ast.Return(
value=ast.Dict(
keys=[ast.Constant(value=n) for n in names],
values=[ast.Name(id=n, ctx=ast.Load()) for n in names],
)
)
)
elif target_body and isinstance(target_body[-1], ast.Expr):
target_body[-1] = ast.Return(value=target_body[-1].value)
if enclosing_assigned:
# Seed each write-back local from the enclosing frame at block entry
# and flush its final value back at block exit, so the block reads and
# rebinds the enclosing local exactly like inline code. A method macro
# such as ``foreach`` runs the block once per item, so the
# per-invocation seed observes the value left by the prior iteration --
# giving correct accumulation for ``total = total + $``.
seeds = [
ast.parse(f"{n} = _dynamic_lookup({n!r})").body[0]
for n in sorted(enclosing_assigned)
]
flushes = [
ast.parse(f"_dynamic_store({n!r}, {n})").body[0]
for n in sorted(enclosing_assigned)
]
if target_body and isinstance(target_body[-1], ast.Return):
target_body[-1:-1] = flushes
else:
target_body.extend(flushes)
target_body[0:0] = seeds
if freevars:
body = [FreeVarTransformer(freevars, frame).visit(stmt) for stmt in body]
fn_def = cast(
ast.FunctionDef,
ast.parse("def __pyc_macro_block__(*__pyc_rest__):\n pass").body[0],
)
fn_def.args.args = [ast.arg(arg=p) for p in params]
fn_def.body = body or [ast.Pass()]
module = ast.Module(body=[fn_def], type_ignores=[])
ast.fix_missing_locations(module)
block_globals = dict(frame.f_globals)
block_globals.setdefault("_dynamic_lookup", _dynamic_lookup)
block_globals.setdefault("_dynamic_store", _dynamic_store)
local_ns: dict = {}
fname = self.make_sandbox_fname()
# Point tracebacks at the user's original block source instead of the
# synthetic `__pyc_macro_block__` sandbox: the parsed body's line numbers
# are 1-based into the (normalized) block, so registering that source in
# linecache lines failures up with what the user actually wrote.
self._register_block_linecache(fname, block_src, param_src)
# exec_raw with instrument=False reuses the active trace (no reset) so
# the already-instrumented nested macros still fire.
self.exec_raw(
module,
global_env=block_globals,
local_env=local_ns,
filename=fname,
instrument=False,
)
return local_ns["__pyc_macro_block__"]
#: Sandbox filenames whose linecache entry holds a block's original source.
#: Tracked so macro sub-lambdas compiled *inside* a block (e.g. ``fork``
#: branches) can alias the same source -- their line numbers already index
#: the block -- giving a traceback that points at the exact failing stage.
_block_linecache_files: set[str] = set()
@classmethod
def _register_block_linecache(
cls, fname: str, block_src: str, param_src: str
) -> None:
"""Register a block's source under its sandbox filename so a traceback
through the compiled block shows the user's lines.
``block_src`` here is already normalized (column-0, no surrounding blank
lines) by :meth:`_compile_block_function`; its lines are 1-for-1 with the
``param_src`` the body was parsed from (placeholder substitution is
in-place), so the friendlier ``$`` form is shown when aligned, falling
back to the desugared ``param_src`` otherwise."""
import linecache
display = block_src
if len(display.split("\n")) != len(param_src.split("\n")):
display = param_src
source = display + "\n"
linecache.cache[fname] = (
len(source),
None,
source.splitlines(keepends=True),
fname,
)
cls._block_linecache_files.add(fname)
# Keep this frame in tracebacks: both pipescript's and ipyflow's
# sandbox-frame filters consult this shared pyccolo registry.
pyc.mark_traceback_visible(fname)
@classmethod
def _alias_block_linecache(cls, frame: FrameType, fn: Any) -> None:
"""If ``fn`` was compiled while executing a block, point its sandbox
filename at that block's source (line numbers already align), so a
traceback through ``fn`` -- e.g. a ``fork`` branch -- shows the exact
stage instead of an empty sandbox frame.
Only the *linecache* (and visibility) are adjusted here; the code object
is left untouched. Renaming an instrumented lambda's code (e.g. to read
``<fork stage>`` rather than ``<lambda>``) breaks pyccolo's runtime
code-identity bookkeeping, so the meaningful source line is the win."""
import linecache
block_fname = frame.f_code.co_filename
if block_fname not in cls._block_linecache_files:
return
code = getattr(fn, "__code__", None)
entry = linecache.cache.get(block_fname)
if code is None or entry is None:
return
linecache.cache.setdefault(code.co_filename, entry)
cls._block_linecache_files.add(code.co_filename)
pyc.mark_traceback_visible(code.co_filename)
def _handle_block_macro(
self, node: ast.Subscript, frame: FrameType, func: str, block_id: int
) -> ast.expr:
__hide_pyccolo_frame__ = True # noqa: F841
from pipescript.tracers.brace_block_tracer import BraceBlockTracer
block_fn = self._compile_block_function(
BraceBlockTracer.block_sources[block_id], frame
)
# Give the block frame a meaningful name in tracebacks (e.g. `map{...}`)
# rather than the synthetic `__pyc_macro_block__`.
try:
block_fn.__code__ = block_fn.__code__.replace(co_name=f"{func}{{...}}")
except Exception:
pass
# expose the compiled function so the macro wrapper can reference it
gname = f"__pyc_block_fn_{id(block_fn)}__"
frame.f_globals[gname] = block_fn
with fast.location_of(node):
ast_lambda: ast.expr = fast.Name(gname, ctx=ast.Load())
if func in self.static_macros and func not in (
"f",
memoize.__name__,
otherwise.__name__,
):
ast_lambda = self._transform_ast_lambda_for_macro(
ast_lambda, func, set()
)
return ast_lambda
def _handle_namespace_block(
self, node: ast.Subscript, frame: FrameType, func: str, block_id: int
) -> object:
__hide_pyccolo_frame__ = True # noqa: F841
from pipescript.tracers.brace_block_tracer import BraceBlockTracer
block_fn = self._compile_block_function(
BraceBlockTracer.block_sources[block_id], frame, namespace=True
)
try:
block_fn.__code__ = block_fn.__code__.replace(co_name=f"{func}{{...}}")
except Exception:
pass
self._alias_block_linecache(frame, block_fn)
try:
namespace = block_fn() # runs the block; returns {name: value, ...}
return self.namespace_block_macros[func](namespace)
except Exception as exc:
# The block body (and the macro builder) run user code; surface its
# exception verbatim instead of letting pyccolo swallow it into a
# None context value (-> cryptic 'NoneType is not a context manager').
self.exc_to_propagate = exc
raise
@pyc.before_subscript_slice(when=is_static_macro, reentrant=True)
def handle_macro(
self, _ret, node: ast.Subscript, frame: FrameType, evt: TraceEvent, *_, **__
):
cached_lambda = self._lambda_cache_lookup(node, frame, evt)
if cached_lambda is not self._not_found:
return cached_lambda
__hide_pyccolo_frame__ = True
func = cast(ast.Name, node.value).id
block_id = block_marker_id(node)
callable_expr: ast.expr
if block_id is not None and func in self.namespace_block_macros:
# Eager: harvest the block's assignments and build the value now, at
# the point the statement runs (so e.g. random init is fresh per run).
# Not memoized in lambda_cache for that reason.
value = self._handle_namespace_block(node, frame, func, block_id)
return lambda: __hide_pyccolo_frame__ and value # noqa: E731
if block_id is not None:
callable_expr = self._handle_block_macro(node, frame, func, block_id)
evaluated_lambda = pyc.eval(callable_expr, frame.f_globals, frame.f_locals)
self._alias_block_linecache(frame, evaluated_lambda)
ret = lambda: __hide_pyccolo_frame__ and evaluated_lambda # noqa: E731
return self._lambda_cache_store(node, frame, evt, ret)
if func in ("macro", "method"):
macro = DynamicMacro.create(node.slice, self, is_method=func == "method")
ret = lambda: __hide_pyccolo_frame__ and macro # noqa: E731
return self._lambda_cache_store(node, frame, evt, ret)
elif func in (fork.__name__, parallel.__name__) and isinstance(
node.slice, ast.Tuple
):
callables: list[ast.expr] = []
max_nargs = 1
expr: ast.expr | None = None
for expr in node.slice.elts:
expr_lambda = self._handle_macro_impl(expr, frame, "f")
callables.append(expr_lambda)
if isinstance(expr_lambda, ast.Lambda):
# Count only placeholder parameters; a branch that closes over
# free variables (e.g. ``A @ $ |> np.sum`` captures ``A``/``np``)
# carries them as *defaulted* params, which the piped value must
# not be expected to fill.
n_placeholders = len(expr_lambda.args.args) - len(
expr_lambda.args.defaults
)
max_nargs = max(max_nargs, n_placeholders)
has_otherwise = (
isinstance(expr, ast.Subscript)
and isinstance(expr.value, ast.Name)
and expr.value.id == otherwise.__name__
)
with fast.location_of(node.slice):
args = [
f"_{arg_ctr}"
for arg_ctr in range(
self.arg_replacer.arg_ctr, self.arg_replacer.arg_ctr + max_nargs
)
]
arg_str = ", ".join(args)
self.arg_replacer.arg_ctr += max_nargs
ast_lambda = cast(
ast.Lambda,
cast(ast.Expr, fast.parse(f"lambda {arg_str}: None").body[0]).value,
)
load = ast.Load()
tuple_elts: list[ast.expr] = []
if has_otherwise:
tuple_elts.append(fast.Constant(value=True))
for lam in callables:
tuple_elts.append(lam)
ast_lambda.body = fast.Call(
func=fast.Name(id=func, ctx=load),
args=[fast.Tuple(tuple_elts, ctx=load)]
+ [fast.Name(arg, ctx=load) for arg in args],
keywords=[],
)
callable_expr = ast_lambda
elif func in (read.__name__, write.__name__):
rw_lambda = self._handle_read_write_macro(func, node.slice) # type: ignore[arg-type]
callable_expr = cast(ast.expr, rw_lambda)
elif func == ntimes.__name__:
ntimes_lambda = self._handle_ntimes_macro(frame, node.slice) # type: ignore[arg-type]
callable_expr = cast(ast.expr, ntimes_lambda)
elif func == once.__name__:
once_call = self._handle_once_macro(frame, node.slice)
callable_expr = cast(ast.expr, once_call)
else:
callable_expr = self._handle_macro_impl(node.slice, frame, func)
evaluated_lambda = pyc.eval(callable_expr, frame.f_globals, frame.f_locals)
self._alias_block_linecache(frame, evaluated_lambda)
ret = lambda: __hide_pyccolo_frame__ and evaluated_lambda # noqa: E731
return self._lambda_cache_store(node, frame, evt, ret)