transformers

AST Transformation Helpers.

Provides logic to reconstruct CST nodes for Infix operators, Inline Lambdas, and Structured Index Selection.

Functions

apply_index_select(→ libcst.Subscript)

Wraps an expression node with a subscript access for a specific integer index.

rewrite_as_inline_lambda(→ libcst.Call)

Wraps arguments in an Immediately Invoked Lambda Expression (IIFE).

rewrite_as_macro(→ libcst.CSTNode)

Replaces an operation call with a python expression defined in the template.

rewrite_as_infix(→ Union[libcst.BinaryOperation, ...)

Transforms a functional call into an infix (binary) or prefix (unary) expression.

Module Contents

transformers.apply_index_select(inner_node: libcst.CSTNode, index: int) → libcst.Subscript

Wraps an expression node with a subscript access for a specific integer index. Safe, structured alternative to string output adapters for tuple destructuring.

Transformation:

Input node: func(…) Output: func(…)[index]

Parameters:

  • inner_node (cst.CSTNode) – The expression node (usually a Call) to slice.

  • index (int) – The integer index to access.

Returns:

The wrapped node.

Return type:

cst.Subscript

transformers.rewrite_as_inline_lambda(lambda_str: str, args: list[libcst.Arg]) → libcst.Call

Wraps arguments in an Immediately Invoked Lambda Expression (IIFE).

transformers.rewrite_as_macro(template: str, args_list: list[libcst.Arg], std_arg_names: list[str]) → libcst.CSTNode

Replaces an operation call with a python expression defined in the template.

Arguments are substituted into the template string using placeholders matching the standard argument names (e.g. {x}).

Parameters:

  • template (str) – The macro string (e.g. “{x} * jax.nn.sigmoid({x})”).

  • args_list (list[cst.Arg]) – The normalized argument nodes for this call.

  • std_arg_names (list[str]) – The names of standard arguments in order.

Returns:

The constructed expression logic.

Return type:

cst.CSTNode

Raises:

  • ValueError – If arguments required by the template are missing.

  • cst.ParserSyntaxError – If the resulting string is invalid Python.

transformers.rewrite_as_infix(_original_node: libcst.Call, args: List[libcst.Arg], op_symbol: str, std_args: List[str]) → libcst.BinaryOperation | libcst.UnaryOperation

Transforms a functional call into an infix (binary) or prefix (unary) expression.