Source code for flowrep.compiler.function

from __future__ import annotations

import dataclasses
import inspect
import typing
from typing import Any

from flowrep import base_models
from flowrep.compiler import annotate, flow_control, statements
from flowrep.parsers import label_helpers
from flowrep.prospective import (
    for_recipe,
    if_recipe,
    try_recipe,
    union_types,
    while_recipe,
    workflow_recipe,
)
from flowrep.retrospective import datastructures


[docs] class NameAllocator: """Mints unique, valid Python identifiers within a single function namespace.""" def __init__(self) -> None: self._used: set[str] = set()
[docs] def fresh(self, hint: str) -> str: base = hint if base_models.is_valid_label(hint) else "_v" candidate = base if base not in self._used else None if candidate is None: candidate = label_helpers.unique_suffix(base, self._used) self._used.add(candidate) return candidate
[docs] def reserve(self, name: str) -> str: self._used.add(name) return name
[docs] @dataclasses.dataclass class Emitter: namespace: dict[str, Any] = dataclasses.field(default_factory=dict) nested_defs: list[str] = dataclasses.field(default_factory=list) module_imports: set[str] = dataclasses.field(default_factory=set) # Module-scope names: nested def names, injected namespace symbols, and # reserved import bindings. Distinct from the per-function local allocator. module_names: NameAllocator = dataclasses.field(default_factory=NameAllocator) workflow_decorator: tuple[str, str] = ("not", "useful") @property def decorator_string(self): mod, qualname = self.workflow_decorator return f"@{mod}.{qualname}"
[docs] @dataclasses.dataclass class FunctionBuilder: name: str decorator: str params: list[str] = dataclasses.field(default_factory=list) body_lines: list[str] = dataclasses.field(default_factory=list) return_annotation: str | None = None return_symbols: list[str] = dataclasses.field(default_factory=list) output_labels: list[str] = dataclasses.field(default_factory=list)
[docs] def render(self) -> str: sig = ", ".join(self.params) # Every generated function is decorated so that exec'ing the source attaches # a `.flowrep_recipe`; this also validates the body parses as a workflow. # Output port names are pinned via decorator args; the parser resolves # `output_labels` ahead of annotation labels and returned-symbol names, so # this round-trips port names without polluting the return annotation. if self.output_labels: label_args = ", ".join(f'"{label}"' for label in self.output_labels) header = f"{self.decorator}({label_args})\n" else: header = f"{self.decorator}\n" header += f"def {self.name}({sig})" if self.return_annotation is not None: header += f" -> {self.return_annotation}" header += ":" ret = "" if self.return_symbols: ret = " return " + ", ".join(self.return_symbols) + "\n" # Only emit `pass` when the body is entirely empty (no lines, no return); # a passthrough workflow has no body lines but does have a return, so # `pass` would be both redundant and illegal in a @flowrep.workflow function. body = self.body_lines or ([] if ret else ["pass"]) indented = "".join(f" {line}\n" for line in body) return header + "\n" + indented + ret
[docs] def referenced_top_level_bindings( node: union_types.RecipeDiscrimination, ) -> typing.Iterator[str]: """Yield the top-level module binding of every import the recipe will emit. Mirrors the structure emission walks: referenced nodes (atomic or referenced workflow) emit a dotted call and are not recursed; reference-free workflows and flow controls are recursed via ``.nodes`` / ``.prospective_nodes``. Builtin exception types are skipped, matching ``_exception_name``. """ reference = getattr(node, "reference", None) if reference is not None: yield reference.info.module.split(".")[0] return if isinstance(node, workflow_recipe.WorkflowRecipe): for child in node.nodes.values(): yield from referenced_top_level_bindings(child) elif isinstance(node, flow_control.FLOW_CONTROL_TYPES): for child in node.prospective_nodes.values(): yield from referenced_top_level_bindings(child) if isinstance(node, try_recipe.TryRecipe): for case in node.exception_cases: for info in case.exceptions: if info.module != "builtins": yield info.module.split(".")[0]
def _inlined_pinned_symbols(recipe: workflow_recipe.WorkflowRecipe) -> set[str]: """Every name the emitter is *forced* to use as a symbol in this function's scope. Python's flow control introduces no scope, so every inlined ``for``/``if``/``while``/ ``try`` body shares one flat namespace with the enclosing ``def``. Most symbols in it are the allocator's to choose, but some are pinned: round-tripping requires a symbol to carry a port name verbatim, or re-parsing regenerates a different port. Pinned names bypass the allocator, so it must learn them *up front* -- a pin emitted later cannot retroactively protect a name the allocator already minted. The pins are flow-control output ports, for-loop variables, for-body output ports, and the input ports of flow-control nodes (which force the feeding sibling's output symbol to match). Descends through inlined flow-control bodies; stops at reference-free WorkflowRecipe peer nodes (emitted as nested defs, with their own scope) and referenced nodes. """ names: set[str] = set(statements.flow_control_input_requirements(recipe).values()) for node in recipe.nodes.values(): _collect_pinned_symbols(node, names) return names def _collect_pinned_symbols( node: union_types.RecipeDiscrimination, names: set[str] ) -> None: if not isinstance(node, flow_control.FLOW_CONTROL_TYPES): return # atomic / referenced / reference-free workflow peer: not inlined here # Every flow-control output port is named for an enclosing symbol, and # `emit_flow_control` pins it to that name. names.update(node.outputs) if isinstance(node, for_recipe.ForEachRecipe): names.update(node.iterated_ports) # `_emit_for_each` pins every body output to its port name. names.update(node.body_node.recipe.outputs) _collect_pinned_symbols_from_body(node.body_node.recipe, names) elif isinstance(node, if_recipe.IfRecipe): for body_case in node.cases: _collect_pinned_symbols_from_body(body_case.body.recipe, names) if node.else_case is not None: _collect_pinned_symbols_from_body(node.else_case.recipe, names) elif isinstance(node, try_recipe.TryRecipe): _collect_pinned_symbols_from_body(node.try_node.recipe, names) for exception_case in node.exception_cases: _collect_pinned_symbols_from_body(exception_case.body.recipe, names) elif isinstance(node, while_recipe.WhileRecipe): _collect_pinned_symbols_from_body(node.case.body.recipe, names) def _collect_pinned_symbols_from_body( body: union_types.RecipeDiscrimination, names: set[str] ) -> None: """Recurse into an inlined body: a WorkflowRecipe's own nodes, or a nested flow-control node. Both are emitted into the current scope.""" if isinstance(body, workflow_recipe.WorkflowRecipe): names.update(statements.flow_control_input_requirements(body).values()) for node in body.nodes.values(): _collect_pinned_symbols(node, names) elif isinstance(body, flow_control.FLOW_CONTROL_TYPES): _collect_pinned_symbols(body, names) # single atomic body: nothing pinned beyond what the caller already reserved
[docs] def emit_nested_workflow_node( node: workflow_recipe.WorkflowRecipe, label: str, emitter: Emitter, alloc: NameAllocator, ) -> str: """Emit a nested @flowrep.workflow-decorated def for a reference-free workflow node. Returns the local function name (the call path for the enclosing body). The nested function is appended to emitter.nested_defs. The function is named after the label's base (pre-suffix) form so that workflow_parser.parse_workflow re-derives the same node label on round-trip via label_helpers.unique_suffix. """ # Use the base of the label (strip _N suffix) so that re-parsing via # unique_suffix reconstructs the same original label. base = statements.label_base(label) # Module-scope: nested defs are emitted at module level, so their names must # be unique across the whole module (not just the parent's locals). Reserve # the chosen name in the parent allocator too, so the parent can't later mint # a local symbol that shadows the def (it is called by bare name in the body). fn_name = emitter.module_names.fresh(base) alloc.reserve(fn_name) builder = emit_workflow_function(node, fn_name, emitter, signature=None) # The @flowrep.workflow decorator is emitted by FunctionBuilder.render itself. rendered = builder.render() if fn_name != base: # Node labels are re-derived from the resolved function's __name__ # (atomic_parser / unique_suffix). When a unique module binding name must # differ from the base -- restore __name__ so re-parsing reconstructs the # original label. Appended to the nested def's text; nested defs are emitted # at module level before the enclosing function, so the assignment executes # before the enclosing @workflow decorator parses the body. rendered += f"{fn_name}.__name__ = {base!r}\n" emitter.nested_defs.append(rendered) return fn_name
def _render_params( recipe: workflow_recipe.WorkflowRecipe, signature: inspect.Signature | None, emitter: Emitter, ) -> list[str]: if signature is None: return list(recipe.inputs) params: list[str] = [] seen_positional_only = False emitted_keyword_marker = False need_pos_only_marker = any( signature.parameters.get(p) and signature.parameters[p].kind == inspect.Parameter.POSITIONAL_ONLY for p in recipe.inputs ) for port in recipe.inputs: param = signature.parameters.get(port) kind = param.kind if param else inspect.Parameter.POSITIONAL_OR_KEYWORD if kind == inspect.Parameter.POSITIONAL_ONLY: seen_positional_only = True else: if seen_positional_only and need_pos_only_marker: params.append("/") need_pos_only_marker = False # only insert once if kind == inspect.Parameter.KEYWORD_ONLY and not emitted_keyword_marker: params.append("*") emitted_keyword_marker = True piece = port has_annotation = ( param is not None and param.annotation is not inspect.Parameter.empty ) has_default = param is not None and param.default is not inspect.Parameter.empty if has_annotation: annotation = typing.cast(inspect.Parameter, param).annotation inlined = annotate.render_annotation(annotation, emitter.module_imports) if inlined is not None: piece += f": {inlined}" else: annotation_name = emitter.module_names.fresh(f"_ann_{port}") emitter.namespace[annotation_name] = annotation piece += f": {annotation_name}" if has_default: default = typing.cast(inspect.Parameter, param).default inlined_default = annotate.render_default(default) if inlined_default is not None: rhs = inlined_default else: default_name = emitter.module_names.fresh(f"_default_{port}") emitter.namespace[default_name] = default rhs = default_name # PEP 8 / black: spaces around '=' only when the parameter is annotated. piece += f" = {rhs}" if has_annotation else f"={rhs}" params.append(piece) # Edge case: the last param was positional-only and no marker emitted yet. if need_pos_only_marker and seen_positional_only and "/" not in params: params.append("/") return params
[docs] def build_signature( inputs: datastructures.InputDataPorts, outputs: datastructures.OutputDataPorts, ) -> inspect.Signature: params = [] for name, port in inputs.items(): default = ( port.default if port.default is not datastructures.NOT_DATA else inspect.Parameter.empty ) annotation = ( port.annotation if port.annotation is not None else inspect.Parameter.empty ) params.append( inspect.Parameter( name, kind=inspect.Parameter.POSITIONAL_OR_KEYWORD, default=default, annotation=annotation, ) ) return inspect.Signature( parameters=params, return_annotation=_build_return_annotation(outputs), )
def _build_return_annotation(outputs: datastructures.OutputDataPorts) -> Any: """Encode per-output-port types as a single function-level return annotation. A single port yields its own annotation; multiple ports yield a ``tuple[...]``. Returns ``inspect.Signature.empty`` when no port carries a type. """ annotations = [port.annotation for port in outputs.values()] if all(annotation is None for annotation in annotations): return inspect.Signature.empty types = [Any if annotation is None else annotation for annotation in annotations] if len(types) == 1: return types[0] return tuple[tuple(types)] # type: ignore
[docs] def emit_workflow_function( recipe: workflow_recipe.WorkflowRecipe, name: str, emitter: Emitter, signature: inspect.Signature | None = None, ) -> FunctionBuilder: alloc = NameAllocator() in_syms = {port: alloc.reserve(port) for port in recipe.inputs} # Pinned symbols bypass the allocator, so it has to know them before it mints # anything, or a fresh name can land on one and be silently overwritten. for pinned in _inlined_pinned_symbols(recipe): alloc.reserve(pinned) params = _render_params(recipe, signature, emitter) lines, out_syms = statements.emit_workflow_body(recipe, in_syms, {}, emitter, alloc) # Output port names are pinned via the decorator (see FunctionBuilder.render), # so the return annotation is free to carry the user's verbatim type. Bind the # whole annotation as a live object (arbitrary types do not round-trip through # repr) and reference it by name; emit nothing when there is no real annotation. return_annotation = None if ( signature is not None and signature.return_annotation is not inspect.Signature.empty ): inlined_return = annotate.render_annotation( signature.return_annotation, emitter.module_imports ) if inlined_return is not None: return_annotation = inlined_return else: ann_name = emitter.module_names.fresh("_ann_return") emitter.namespace[ann_name] = signature.return_annotation return_annotation = ann_name # Imports are hoisted to the module preamble via emitter.module_imports, so # the body carries only the optional docstring (which must come first so the # parser's skip_docstring recognises it and `description` round-trips) and # the emitted statements. body_lines = lines if recipe.description is not None: body_lines = [repr(recipe.description)] + body_lines builder = FunctionBuilder( name=name, decorator=emitter.decorator_string, params=params, body_lines=body_lines, return_annotation=return_annotation, return_symbols=[out_syms[p] for p in recipe.outputs], output_labels=list(recipe.outputs), ) return builder