CODEWITHSUNDEEP
pythonabstract-syntax-tree
BBrooklyn
BBrooklyn

Reputation: 429

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.

Upvotes: 0

Views: 2098

Answers (1)

wjandrea
wjandrea

Reputation: 33179

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

Upvotes: 6

Related Questions