Search code examples
pythonpython-3.xvariablesabstract-syntax-tree

How to parse python code and only get variables with no indentation?


I'm trying to extract all variables in my code that have no indentation, here is a simple example:

import ast
import astunparse
class AnalysisNodeVisitor(ast.NodeVisitor):
    def __init__(self, nodename):
        super().__init__()
        self.nodename=nodename
        self.setVariables={}
    def visit_Assign(self,node):
        print(self.nodename,astunparse.unparse(node.targets))
        #print(ast.dump(node))

        
Analyzer=AnalysisNodeVisitor("example1")
tree=ast.parse('''\
a = 10

for i in range(10):
    b=x+i

if(x==10):
    c=36

d=20
''')
Analyzer.visit(tree)

output:

example1 a

example1 b

example1 c

example1 d

This example prints out all the 4 variables(a,b,c and d) and I'm wondering if there is a way to make it only print the assignments a,d since there is no indentation before their assignments?

I tried to dump the nodes and see if there is something I can use to filter out the variables but I can't find anything (for example it just outputs this : 'Assign(targets=[Name(id='a', ctx=Store())], value=Num(n=10))').


Solution

  • I think if you just look at tree.body, you can find all your top-level assignment statements.

    Running this code:

    import ast
    import astunparse
    
    tree=ast.parse('''\
    a = 10
    
    for i in range(10):
        b=x+i
    
    if(x==10):
        c=36
    
    d=20
    ''')
    
    for thing in tree.body:
        if isinstance(thing, ast.Assign):
            print(astunparse.unparse(thing.targets))
    

    Produces:

    
    a
    
    
    d