ml_switcheroo.core.rewriter.calls.transformers

AST Transformation Helpers.

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

Functions

apply_output_adapter(→ libcst.Call)	Wraps a node with a lambda adapter to normalize output via string parsing.
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

ml_switcheroo.core.rewriter.calls.transformers.apply_output_adapter(inner_node: libcst.CSTNode, adapter_str: str) → libcst.Call

Wraps a node with a lambda adapter to normalize output via string parsing. Transform: node -> (adapter_str)(node)

ml_switcheroo.core.rewriter.calls.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

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

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

ml_switcheroo.core.rewriter.calls.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.

ml_switcheroo.core.rewriter.calls.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.