Example to insert python code through AST

[hshen14]

Hi,

Is there any example to show how to insert python code automatically through AST?

Thanks.

[berkerpeksag]

Here's a quick and dirty example:

import ast
import astor

def add_new_argument_to_function(node, name):
    node.args.args.append(ast.arg(arg=name, annotation=None))
    for body_node in node.body:
        if isinstance(body_node, ast.Return):
            elts = [
                # We assume ast.Name.
                body_node.value,
                ast.Name(id=name, ctx=ast.Load()),
            ]
            body_node.value = ast.Tuple(elts=elts, ctx=ast.Load())
    return node

name = "foo"
source = f"""
def {name}(a):
    return a
"""

tree = ast.parse(source)

for i, node in enumerate(tree.body):
    if isinstance(node, ast.FunctionDef) and node.name == name:
        tree.body[i] = add_new_argument_to_function(node, "b")

print(astor.to_source(tree))

This will print:

def foo(a, b):
    return a, b

I think it would be nice to add this one and a couple more examples to the documentation. Patches welcome.