Package: ast-tools
Toolbox for working with the Python ASTpip install ast_toolsUseful ReferencesGreen Tree Snakes - the missing Python AST docsPassesast_tools provides a number of passes for rewriting function and classes (could
also work at the module level however no such pass exists). Passes are
applied with theapply_passesdecorator:@apply_passes([pass1(),pass2()])deffoo(...):...Each pass takes as arguments an AST, an environment, and metadata and
returns (possibly) modified versions of each.apply_passesbegins a chain of rewrites by first looking
up the ast of the decorated object and gather attempts to gather locals
and globals from the call site to build the environment.After all rewrites have runapply_passesserializes and
execute the rewritten ast.Know IssuesCollecting the ASTapply_passesrelies oninspect.getsourceto get the
source of the decorated definition (which is then parsed to get the initial ast).
However,inspect.getsourcehas many limitations.Collecting the Environmentapply_passesdoes its best to infer the environment
however there is no way to do this in a fully correct way.  Users are
encouraged to pass environment explicitly:@apply_passes(...,env=SymbolTable(locals(),globals()))deffoo(...):...Wrapping the apply_passes decoratorTheapply_passesdecorator must not be wrapped.As decorators are a part of the AST of the object they are applied to
they must be removed from the rewritten AST before it is executed.  If they
are not removed rewrites will recurse infinitely as@apply_passes([...])deffoo(...):...would becomeexec('''\@apply_passes([...])def rewritten_foo(...): ...''')Note: this would invokeapply_passes([...])onrewritten_fooTo avoid this theapply_passesdecorator filters itself from the decorator list.  If, however,
the decorator is wrapped inside another decorator, this will fail.Inner decorators are called multiple timesDecorators that are applied before a rewrite group will be called multiple times.
Seehttps://github.com/leonardt/ast_tools/issues/46for detailed explanation.
To avoid this users are encouraged to make rewrites the inner most decorators
when possible.MacrosLoop UnrollingUnroll loops using the patternfor<var>inast_tools.macros.unroll(<iter>):...<iter>should be an iterable object that produces integers (e.g.range(8))
that can be evaluated at definition time (can refer to variables in the scope
of the function definition)For example,fromast_tools.passesimportapply_passes,loop_unroll@apply_passes([loop_unroll()])deffoo():foriinast_tools.macros.unroll(range(8)):print(i)is rewritten intodeffoo():print(0)print(1)print(2)print(3)print(4)print(5)print(6)print(7)You can also use a list ofints, here's an example that also uses a reference
to a variable defined in the outer scope:fromast_tools.passesimportapply_passes,loop_unrollj=[1,2,3]@apply_passes([loop_unroll()])deffoo():foriinast_tools.macros.unroll(j):print(i)becomesdeffoo():print(1)print(2)print(3)Inlining If StatementsThis macro allows you to evaluateifstatements at function definition time,
so the resulting rewritten function will have theifstatements marked
"inlined" removed from the final code and replaced with the chosen branch based
on evaluating the condition in the definition's enclosing scope.ifstatements are marked by using the formif inline(...):whereinlineis
imported from theast_tools.macrospackage.ifstatements not matching
this pattern will be ignored by the rewrite logic.Here's an examplefromast_tools.macrosimportinlinefromast_tools.passesimportapply_passes,if_inliney=True@apply_passes([if_inline()])deffoo(x):ifinline(y):returnx+1else:returnx-1importinspectassertinspect.getsource(foo)==f"""\def foo(x):return x + 1"""