ast - Get the list of only the variable names in an expression

1.1k Views Asked by At

I'm using the Python ast module to obtain a list of variable names in a Python expression. For example, the expression [int(s[i:i + 3], 2) for i in range(0, len(s), 3)] should return a singleton list with the variable name s like [s]. I've tried the following code snippet -

names = [
    node.id for node in ast.walk(ast.parse(formula)) 
    if isinstance(node, ast.Name)
]

which returns a list of variables plus function names in the ast -

['int', 'i', 's', 'range', 'i', 'len', 's', 'i']

But I don't want to include function names like range, len, int and the iterator i.

1

There are 1 best solutions below

1
Lyubomyr Ivanitskiy On

You can use similar approach but then filter out names from builtins. Something like this

import ast
import builtins

def get_variables(expression):
    tree = ast.parse(expression)
    variables = []
    for node in ast.walk(tree):
        if isinstance(node, ast.Name):
            variables.append(node.id)
    return tuple(v for v in set(variables) if v not in vars(builtins))