Python ast package: traversing object hierarchies

Question

The following is a Python code snippet using the ast and symtable packages. I am trying to parse the code and check the types. But I don't understand how to traverse objects to get to the actual variable being referenced.

The following code implements a NodeVisitor, and a function is presented to the compiler and parsed by the compiler and the ast walked. The function being analyzed (eval_types) is passed a couple objects.

Below are the code chunks that make up the example. I have added some comments for each chunk. To run the code, the "chunks" need to be reassembled.

The imports and a function to un-indent a block of code for parsing.

import inspect
import ast
import symtable
from tokenize import generate_tokens, untokenize, INDENT
from cStringIO import StringIO

# _dedent borrowed from the myhdl package (www.myhdl.org)
def _dedent(s):
    """Dedent python code string."""

    result = [t[:2] for t in generate_tokens(StringIO(s).readline)]
    # set initial indent to 0 if any
    if result[0][0] == INDENT:
        result[0] = (INDENT, '')
    return untokenize(result)

The following is the node visitor, it has the generic unhandled and name visitor overloads.

class NodeVisitor(ast.NodeVisitor):
    def __init__(self, SymbolTable):
        self.symtable = SymbolTable
        for child in SymbolTable.get_children():
            self.symtable = child
            print(child.get_symbols())

    def _visit_children(self, node):
        """Determine if the node has children and visit"""
        for _, value in ast.iter_fields(node):
            if isinstance(value, list):
                for item in value:
                    if isinstance(item, ast.AST):
                        print('  visit item %s' % (type(item).__name__))
                        self.visit(item)

            elif isinstance(value, ast.AST):
                print('  visit value %s' % (type(value).__name__))
                self.visit(value)

    def generic_visit(self, node):
        print(type(node).__name__)
        self._visit_children(node)

    def visit_Name(self, node):
        print('  variable %s type %s' % (node.id,
                                         self.symtable.lookup(node.id)))
        print(dir(self.symtable.lookup(node.id)))

The following are some simple classes that will be used in the function that will be parsed and analyzed with the AST.

class MyObj(object):
    def __init__(self):
        self.val = None

class MyObjFloat(object):
    def __init__(self):
        self.x = 1.

class MyObjInt(object):
    def __init__(self):
        self.x = 1

class MyObjObj(object):
    def __init__(self):
        self.xi = MyObjInt()
        self.xf = MyObjFloat()

The following is the test function, the eval_types function is the function that will be analyzed with the AST.

def testFunc(x,y,xo,z):

    def eval_types():
        z.val = x + y + xo.xi.x + xo.xf.x

    return eval_types

The code to execute the example, compile the function and analyze.

if __name__ == '__main__':
    z = MyObj()
    print(z.val)
    f = testFunc(1, 2, MyObjObj(), z)
    f()
    print(z.val)
    s = inspect.getsource(f)
    s = _dedent(s)
    print(type(s))
    print(s)

    SymbolTable = symtable.symtable(s,'string','exec')
    tree = ast.parse(s)
    v = NodeVisitor(SymbolTable)
    v.visit(tree)

The following is an example output up to the first name visit.

Module
  visit item FunctionDef
FunctionDef
  visit value arguments
arguments
  visit item Assign
Assign
  visit item Attribute
Attribute
  visit value Name
  variable z type <symbol 'z'>
['_Symbol__flags', '_Symbol__name', '_Symbol__namespaces', 
 '_Symbol__scope', '__class__', '__delattr__', '__dict__', 
 '__doc__', '__format__', '__getattribute__', '__hash__', 
 '__init__', '__module__', '__new__', '__reduce__', '__reduce_ex__', 
 '__repr__', '__setattr__', '__sizeof__', '__str__', 
 '__subclasshook__', '__weakref__', 'get_name', 'get_namespace', 
 'get_namespaces', 'is_assigned', 'is_declared_global', 
 'is_free', 'is_global', 'is_imported', 'is_local', 
 'is_namespace', 'is_parameter', 'is_referenced']

Creating the node visitor doesn't seem to bad but I can't figure out how to traverse an object hierarchy. In the general case the variable being accessed could be buried deep in a object. How to get to the actual variable being accessed from the ast visitor? I only see that an object is at the node but no additional information what the result variable access is.

Answer 1

I don't know if you're still looking for this, but it looks like you just need to add a visit_Attribute and traverse backwards. If you add this to your example:

def visit_Attribute(self, node):
    print('  attribute %s' % node.attr)
    self._visit_children(node)

Then the output for xo.xf.x is:

Add
  visit value Attribute
  attribute x
  visit value Attribute
  attribute xf
  visit value Name
  variable xo type <symbol 'xo'>
  visit value Load

Depending what you want to do with this, you would just need to store the attributes in a list until a Name is encountered, then reverse them.