Search code examples
pythonabstract-syntax-tree

Using Python AST to determine if a string has been hard-coded into a print call


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.


Solution

  • 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