0

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.

wjandrea
  • 28,235
  • 9
  • 60
  • 81
BBrooklyn
  • 350
  • 1
  • 3
  • 15
  • Tip: `ast.dump` is great for helping you figure out the structure of an individual AST node. Try using `ast.dump(nodes[2])`. – SuperStormer Feb 17 '22 at 01:43
  • I've tried that and I would rather get the answer through ast objects so as to be more robust, rather than parse a string representation of the the ast – BBrooklyn Feb 17 '22 at 01:46
  • 1
    @BBrooklyn SuperStormer is saying you can use a dump to help *you* understand the structure, i.e. see that a `Call` has `args`, one of which is a `Constant` with a `value`. – wjandrea Feb 17 '22 at 02:53

1 Answers1

4

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
wjandrea
  • 28,235
  • 9
  • 60
  • 81
  • Thanks, this is very helpful. Is there a more direct way of programmatically accessing a function call's passed in arguments? It seems so clunky to have to do all these isinstance checks, even after the Call node has been found. – BBrooklyn Feb 17 '22 at 18:50
  • @BBrooklyn Not that I'm aware of, though I don't use AST very much. You have to determine, for example, that `arg` is a `Constant` before you can check its `value`. – wjandrea Feb 17 '22 at 19:45
  • @BBrooklyn Oh, actually, you could use structural pattern matching. For example: `case ast.Call(func=ast.Name(id='print'), args=[ast.Constant(value="hard coded")]):`. But I'm not sure how to translate the `any`; this checks that there's exactly one positional argument. – wjandrea Feb 17 '22 at 20:01
  • @BBrooklyn FYI, I added structural pattern matching to the answer, with a translation of the `any`. Turns out [there's no way to translate it all in one pattern](/a/67796857/4518341). – wjandrea Feb 17 '22 at 22:58