How to work with python ast module to analyse if-statements

Question

I have to analyse python code that contains if statements and I found the ast module: https://docs.python.org/3.8/library/ast.html Somehow the documentation is not self-explanatory. I found an example here: https://www.mattlayman.com/blog/2018/decipher-python-ast/ that uses the ast.NodeVisitor helper class but I am struggeling how to adopt this example to get the details of an if statement.

Code to parse:

toggleSwitch = False

# check for someValue in the key-value store
if 'someValue' in context['someKey']:
    toggleSwitch = True

Code of the analyser:

class Analyzer(ast.NodeVisitor):
    def visit_If(self, node):
        print("If:",node.test.left)
        self.stats["if"].append(node.body)
        self.generic_visit(node)

I expect to access the 'someValue' element in some kind of attribute of the node inside the visit_If function, but I don't know how to do it exactly.

Answer 1

GreenTreeSnakes has pretty extensive documentation on the nodes in a Python AST tree.

I don't know if you're actually parsing the code into an ast tree or not, so I'll include that here.

Parse the code into a tree:

code = '''toggleSwitch = False

# check for someValue in the key-value store
if 'someValue' in context['someKey']:
    toggleSwitch = True'''

import ast
tree = ast.parse(code)

Then in your Analyzer class, you can get the someValue symbol from the s attribute of an _ast.Str node.

class Analyzer(ast.NodeVisitor):
    def __init__(self):
        self.stats = {'if': []}

    def visit_If(self, node):
        # Add the "s" attribute access here
        print("If:", node.test.left.s)
        self.stats["if"].append(node.body)
        self.generic_visit(node)

    def report(self):
        pprint(self.stats)

>>> a = Analyzer()
>>> a.visit(tree)
If: someValue

For an If node, the attributes go test (_ast.Compare) → left (_ast.Str) → s (str).