Source code for meta.asttools.visitors.symbol_visitor
'''
Created on Aug 3, 2011
@author: sean
'''
from meta.asttools.visitors import Visitor
import ast
class SymbolVisitor(Visitor):
def __init__(self, ctx_types=(ast.Load, ast.Store)):
if not isinstance(ctx_types, (list, tuple)):
ctx_types = (ctx_types,)
self.ctx_types = tuple(ctx_types)
def visitDefault(self, node):
ids = set()
for child in self.children(node):
if isinstance(child, (tuple, list)):
for item in child:
ids.update(self.visit(item))
elif isinstance(child, ast.AST):
ids.update(self.visit(child))
return ids
def visitName(self, node):
if isinstance(node.ctx, self.ctx_types):
return {node.id}
else:
return set()
def visitalias(self, node):
name = node.asname if node.asname else node.name
if '.' in name:
name = name.split('.', 1)[0]
if ast.Store in self.ctx_types:
return {name}
else:
return set()
[docs]def get_symbols(node, ctx_types=(ast.Load, ast.Store)):
'''
Returns all symbols defined in an ast node.
if ctx_types is given, then restrict the symbols to ones with that context.
:param node: ast node
:param ctx_types: type or tuple of types that may be found assigned to the `ctx` attribute of
an ast Name node.
'''
gen = SymbolVisitor(ctx_types)
return gen.visit(node)