Skip to content

Writing transforms

A transform receives the whole recovered tree as an ast.Module whose single statement is your function's FunctionDef, and must return the Module. Two forms are accepted:

@rewrite(MyTransformer())            # an ast.NodeTransformer instance
@rewrite(lambda tree: mutate(tree))  # any Module -> Module callable

The tree arrives with real file positions: line numbers are offset to the function's true location before your transform runs, so anything you keep carries correct locations through the rewrite, and anything you copy locations from is already right.

Locations: keep the user's, fill your own

The rule that survived our traceback-fidelity measurements: splice user code into your templates, never your template code around theirs without locations. Concretely:

  • Nodes you preserve keep their own locations. Do nothing.
  • Nodes you build fresh: either ast.copy_location(new, anchor) from a node they replace, or leave them location-free and let redef fill them. After your transform returns, ast.fix_missing_locations runs once; it fills only absent locations, inheriting from parents, and never touches the correct ones the user's nodes carry.

A failure inside rewritten code then reports the line the user wrote. CPython 3.11+ carets (PEP 657) follow the same locations, so columns come along.

class WrapEach(ast.NodeTransformer):
    """Wrap every statement in a try/except; user lines still report correctly."""
    def visit_FunctionDef(self, node):
        node.body = [
            ast.Try(body=[stmt], handlers=[HANDLER], orelse=[], finalbody=[])
            for stmt in node.body      # stmt keeps its authored location
        ]
        return node

Hygiene: fresh names cost one call

A rewriter that introduces temporaries can silently overwrite a binding the user holds, and nothing reports it. gensym scans the tree for every Name and argument and returns an identifier that collides with none of them:

from redef import gensym

def visit_FunctionDef(self, node):
    tmp = gensym("captured", avoid=node)
    ...

This solves the capture half of hygiene. The other half, making names your generated code refers to resolve where you defined them, does not arise here: redef has no macro definition environment, and your generated code resolves in the rewritten function's own module. If you inject references to your library, inject them through the module namespace (_helper = helper at module top in the user's file, or attribute access on an imported module).

What a transform cannot rewrite

  • Default arguments. The original default objects are copied onto the new function, preserving identity semantics for mutable defaults. Your transform sees the default expressions in the tree, and changes to them are discarded.
  • The signature's identity metadata. __qualname__, __defaults__, __kwdefaults__, and __dict__ come from the original.
  • Anything outside the function. The unit is one def. If you need whole-module rewriting, you want an import hook, which is a different tool with different costs; see the research.

Displaying expansions

rewrite.expansion(f) returns the unparsed source of what actually ran through compile. Print it in your own tooling, show it in error messages, snapshot it in tests. It is the single best debugging affordance a rewriter can offer its users; we borrowed the idea from mcpyrate's stepwise expansion display.