CODEWITHSUNDEEP
pythonabstract-syntax-tree
Houy Narun
Houy Narun

Reputation: 1725

How to get actual user defined variable name in Python Abstract Syntax Tree (AST)?

From the given code snippet below, I tried to extract the actual user defined variable name other than the python built-in variable in order to check them agaist my naming convention rules.

ast_example.py

import ast
from pprint import pprint


def main():
    FirstName = "Johnny"
    LastName = "David"
    Variable_One = "Variable1"
    Variable_Two = "Variable2"
    Variable_Three = "Variable3"
    with open("ast_example.py", "r") as source:
        tree = ast.parse(source.read())

    analyzer = Analyzer()
    analyzer.visit(tree)
    analyzer.report()


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

    def visit_Name(self, node):
        print "Node: ",node.id
        self.stats["variable"].append(node.id)
        self.generic_visit(node)

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


if __name__ == "__main__":
    main()

However, after I executed the above code snippet, it resulted in not only the variables I wanted but also every python built in variables as well such as self, __name__, open, source, etc which are those I want to exclude.

{'variable': ['FirstName',
              'LastName',
              'Variable_One',
              'Variable_Two',
              'Variable_Three',
              'open',
              'source',
              'tree',
              'ast',
              'source',
              'analyzer',
              'Analyzer',
              'analyzer',
              'tree',
              'analyzer',
              'ast',
              'self',
              'self',
              'self',
              'node',
              'self',
              'node',
              'self',
              'node',
              'self',
              'pprint',
              'self',
              '__name__',
              'main']}

How can I exclude those built in variables? Thanks

Upvotes: 2

Views: 1121

Answers (1)

BlueSheepToken
BlueSheepToken

Reputation: 6169

Instead of visiting Name you might want to visist Assign which is the node that represents assignement.

So the code looks something like this :

def visit_Assign(self, node):
    for target in node.targets:
        self.stats["variable"].append(target.id)
    self.generic_visit(node)

Here targets, represents the multiple values for assignement such as : a, b = 0, 1

ref: python ast doc

Upvotes: 2

Related Questions