def visitAssign(self, node):
assign_nodes = node.nodes
if isinstance(assign_nodes[0], AssTuple):
assigns = assign_nodes[0].nodes
else:
assigns = assign_nodes
assign_names = [n.name for n in assigns]
expression = node.expr
# Collect names.
nf = NameFinder()
walk(expression, nf)
self.names.update(nf.names)
self.expressions.append((assign_names, expression))
def count_function_args(ast):
""" Given an ast node, finds the total number of arguments in FunctionCalls.
This is probably only useful if you know the node to have only one
function call.
"""
return walk(ast, FunctionArgWalker()).num_args
def find_local_defs(ast):
""" Given Block's ast, return a dictionary of def names and
the Function ast that represents them. This will be used
to create a LocalFunctionInfo.
"""
res = walk(ast, LocalDefinitionsWalker()).local_funcs
return res
def from_ast(cls, ast):
""" Create an ExecutionModel object from the ast. """
# Build dictionaries for import information and local defs
info = find_functions(ast)
statement_info = walk(ast, StatementWalker())
statements = parse_stmts_info(statement_info,info)
return cls(statements=statements)
def visitWhile(self,node):
""" Explore while loops
"""
gfunc_def = node.test
sub_statement_info = walk(node.body, StatementWalker())
group_struct = {'type':'while',
'stmts':sub_statement_info,
'gfunc':gfunc_def}
self.groups.append(group_struct)
def visitFor(self,node):
""" Explore for loops
"""
gfunc_def = (node.assign, node.list)
sub_statement_info = walk(node.body, StatementWalker())
group_struct = {'type':'for',
'stmts':sub_statement_info,
'gfunc':gfunc_def}
self.groups.append(group_struct)
def function_inputs_from_call_ast(ast, info=None):
""" Returns function names and input/output binding
information from the CallFunc.
"""
func = walk(ast, FunctionCallWalker())
func_names = func.func_names
func_args = func.func_args
if len(func_names) == 1:
return func_names, func_args
else:
return None
def test_complex_statement_walker():
code = "\n".join([
"from math import sin",
"x = 10",
"y = 10 + x",
"z = sin(y) + b",
])
walker = walk(parse(code), StatementWalker())
assert len(walker.names) == 4
assert set(walker.names) == set(['x', 'y', 'b', 'sin'])
assert len(walker.imported_names) == 1
assert set(walker.imported_names) == set(['sin'])
assert len(walker.expressions) == 3
def test_simple_statement_walker():
const_code = "a = 10"
walker = walk(parse(const_code), StatementWalker())
assert len(walker.names) == 0
assert len(walker.imported_names) == 0
assert len(walker.expressions) == 1
const_code = "a, b = 10"
walker = walk(parse(const_code), StatementWalker())
assert len(walker.names) == 0
assert len(walker.imported_names) == 0
assert len(walker.expressions) == 1
const_code = "a = 10; b = 20"
walker = walk(parse(const_code), StatementWalker())
assert len(walker.names) == 0
assert len(walker.imported_names) == 0
assert len(walker.expressions) == 2
const_code = "a = 10 + 20"
walker = walk(parse(const_code), StatementWalker())
assert len(walker.names) == 0
assert len(walker.imported_names) == 0
assert len(walker.expressions) == 1
def function_returns_from_ast(ast):
""" Find the function returns from a function's ast.
Parameters:
-----------
code: Str
source code of the function
"""
result = None
if ast is not None:
return_vals = walk(ast, ReturnVariablesFinder()).return_vals
if len(return_vals) != 0:
# The last set of return vals is likely to be the most used.
result = return_vals[-1]
else:
result = []
return result
def visitCallFunc(self, node):
if isinstance(node.node, Getattr):
getattr_node = node.node
decorated_name = ""
while (isinstance(getattr_node.expr, Getattr)):
decorated_name = getattr_node.attrname + "." + decorated_name
getattr_node = getattr_node.expr
decorated_name = "%s.%s.%s" % (getattr_node.expr.name,
getattr_node.attrname, decorated_name)
decorated_name = decorated_name[:-1]
self.function_names.append(decorated_name)
else:
self.function_names.append(node.node.name)
# fixme: The fact that I have to do this means that something is very
# wrong somewhere, but where? walk is suppposed to walk the whole
# AST tree! Why isn't this happening?
for arg in node.args:
self.function_names.extend(walk(arg, BlockInfoWalker()).function_names)
def _retrieve_inputs_and_outputs(code = "", reserved_inputs=None, \
reserved_outputs=None):
""" Parse code to retrieve inputs and outputs
taking into account who are the reserved_inputs and
reserved_outputs.
This function returns list of InputVariable, OutputVariable.
"""
if code =="":
return
ast = parse(code)
walker = walk(ast, StatementWalker())
imported_names = walker.imported_names
outputs = []
notinputs = set()
for assigns, _ in walker.expressions:
for name in assigns:
if reserved_outputs is None or name not in reserved_outputs:
outputs.append(OutputVariable(name=name, binding=name))
notinputs.add(name)
notinputs.update(imported_names)
# Add function definition to the list of notinputs
function_calls = [expr for assigns, expr in walker.expressions \
if isinstance(expr, CallFunc) ]
[notinputs.add(call.node.name) for call in function_calls]
# Add the builtins, too.
notinputs.update(builtin_names)
inputs = []
for name in walker.names:
if name not in notinputs:
if reserved_inputs is None or name not in reserved_inputs:
inputs.append(InputVariable(name=name, binding=name))
return inputs,outputs
def retrieve_inputs_and_outputs(code="",reserved_inputs=[],reserved_outputs=[]):
""" Parse code to retrieve inputs and outputs
taking into account who are reserved_inputs and
reserved_outputs.
FIXME: fix the StatementWalker object to manage function_call
without returned value (no assignment).
"""
ast = parse(code)
walker = walk(ast, StatementWalker())
imported_names = walker.imported_names
outputs = []
notinputs = set()
for assigns, _ in walker.expressions:
for name in assigns:
if reserved_outputs is None or name not in reserved_outputs:
outputs.append(name)
notinputs.add(name)
notinputs.update(imported_names)
# Add function definition to the list of notinputs
function_calls = [expr for assigns, expr in walker.expressions \
if isinstance(expr, CallFunc) ]
[notinputs.add(call.node.name) for call in function_calls]
# Add the builtins, too.
notinputs.update(builtin_names)
inputs = []
for name in walker.names:
if name not in notinputs:
if reserved_inputs is None or name not in reserved_inputs:
# Icky. We don't know that the user gave code in the correct
# order.
inputs.append(name)
return inputs,outputs
def _update_code(self, code):
""" Update the expression code.
"""
ast = parse(code)
walker = walk(ast, StatementWalker())
self.imported_names = walker.imported_names
outputs = []
notinputs = set()
for assigns, expr in walker.expressions:
for name in assigns:
outputs.append(OutputVariable(name=name, binding=name))
notinputs.add(name)
notinputs.update(self.imported_names)
# Add the builtins, too.
notinputs.update(builtin_names)
inputs = []
for name in walker.names:
if name not in notinputs:
# Icky. We don't know that the user gave code in the correct
# order.
inputs.append(InputVariable(name=name, binding=name))
self.inputs = inputs
self.outputs = outputs
def function_names(block):
""" Return all the function names within a block.
"""
return walk(block.ast, BlockInfoWalker()).function_names
def find_functions(ast):
""" Given Block's ast, return a dictionary of imports and local
functions. Dictionary is of function name -> Callable.
"""
res = walk(ast, FunctionInfoWalker())
return res.function_callables
def rename_variable(ast, old, new):
""" Replace all references to 'old' with 'new'.
"""
return walk(ast, BlockNameReplacer(old, new))
