using ast module to transform random constants in Python source

Question:

I’m interested in writing a program that uses Python’s built-in AST module to randomly modify constants in arbitrary Python source.

This transformation will likely involve a traversal of the abstract syntax tree representation using operations defined by the AST module. The module offers two options: first, ast.walk() returns references to all nodes in the AST, but does not offer any contextual information, making reassembling the tree impossible. Secondly, the documentation describes a second method involving the ast.NodeTransformer class: several sources of documentation describe generally how to use NodeTransformer.

However, the NodeTransformer documentation fails to mention how to randomly apply conditional substitutions to the AST. Specifically, I would like to modify this feature to create a function that selects a random node in the ast, chooses a constant associated with the node at random, and replace that constant with a randomly selected constant of the same type.

I suspect that I’m struggling to understand how to properly modify NodeTransformer because I rarely program in an object oriented style (usually adhering to the functional paradigm). Hopefully pointing me in the right direction will come easily to one of you.

Asked By: David Shaked

||

Answers:

If all you want to do is modify constants randomly, then you don’t really need node context. You can just walk the tree looking for constant nodes, and if you see one, change its value. Here’s a simple example:

source = """
x = 2
"""

# exec the source as-is
mod = {}
exec compile(source, '<blah>', 'exec') in mod

print(mod['x'])
# prints 2

t = ast.parse(source)
# change all numerical constants to 8
for node in ast.walk(t):
    if isinstance(node, ast.Num):
        node.n = 8
# exec the modified AST
modMod = {}
exec compile(t, '<blah>', 'exec') in modMod

print(modMod['x'])
# prints 8

If you want to make a random choice about whether to modify a constant, you could do that as well. I’m not sure I understand your problem statement about choosing “a random node in the AST”. The hierarchical nature of the AST means that a random node could be anything from a single constant to the entire module, so it seems a bit odd to first pick a random node and then pick a constant node within that. Why not just pick a random constant node in the first place?

Answered By: BrenBarn