I'm trying to write a function that will use AST to determine if a certain string has been hard-coded into a print function call. For example if this is the entire Python file:
print("hard coded")
I get the tree as follows:
with open(filename, 'r') as f:
tree = ast.parse(f.read())
nodes = [node for node in ast.walk(tree)]
Now nodes
contains [<_ast.Module object at 0x7fa250286150>, <_ast.Expr object at 0x7fa250286050>, <_ast.Call object at 0x7fa2502642d0>, <_ast.Name object at 0x7fa228014610>, <_ast.Str object at 0x7fa228014710>, <_ast.Load object at 0x7fa2280080d0>]
By testing attrs I can find that nodes[2]
has a .func.id
and nodes[2].func.id == 'print'
, so there's my print command. But how do I find the argument passed to this print command? The hard coded string is in the ast.Str
object that appears at index 4 of the list, but I need to establish specifically that the string was passed to print, not just that the string appears in the file.
A Call
object has an args
attribute you can use, for example:
for node in ast.walk(tree):
if (
isinstance(node, ast.Call) # It's a call
and isinstance(node.func, ast.Name) # It directly invokes a name
and node.func.id == 'print' # That name is `print`
):
# Check if any arg is the one we're looking for
print(any(
arg.value == "hard coded"
for arg in node.args
if isinstance(arg, ast.Constant)
))
Output:
True
In Python 3.10+, you could use structural pattern matching to simplify the isinstance
stuff:
for node in ast.walk(tree):
match node:
# If it's a call to `print`
case ast.Call(func=ast.Name(id='print')):
# Check if any arg is the one we're looking for
for arg in node.args:
match arg:
case ast.Constant(value="hard coded"):
print('found')
break
Output:
found