def _sub_for_func_ast(ast, func_name, func_vars, func_expr_ast):
"""
Return an ast with the function func_name substituted out.
"""
if isinstance(ast, Call) and ast2str(ast.func) == func_name\
and func_vars == '*':
working_ast = copy.deepcopy(func_expr_ast)
new_args = [_sub_for_func_ast(arg_ast, func_name, func_vars,
func_expr_ast) for arg_ast in ast.args]
# This subs out the arguments of the original function.
working_ast.values = new_args
return working_ast
if isinstance(ast, Call) and ast2str(ast.func) == func_name\
and len(ast.args) == len(func_vars):
# If our ast is the function we're looking for, we take the ast
# for the function expression, substitute for its arguments, and
# return
working_ast = copy.deepcopy(func_expr_ast)
mapping = {}
for var_name, arg_ast in zip(func_vars, ast.args):
subbed_arg_ast = _sub_for_func_ast(arg_ast, func_name, func_vars,
func_expr_ast)
mapping[var_name] = subbed_arg_ast
_sub_subtrees_for_vars(working_ast, mapping)
return working_ast
ast = AST.recurse_down_tree(ast, _sub_for_func_ast,
(func_name, func_vars, func_expr_ast,))
return ast
def _sub_subtrees_for_vars(ast, ast_mappings):
"""
For each out_name, in_ast pair in mappings, substitute in_ast for all
occurances of the variable named out_name in ast
"""
if isinstance(ast, Name) and ast2str(ast) in ast_mappings:
return ast_mappings[ast2str(ast)]
ast = AST.recurse_down_tree(ast, _sub_subtrees_for_vars, (ast_mappings,))
return ast
def _extract_vars_ast(ast, vars_found):
"""
Appends the asts of the variables used in ast to vars_found.
"""
if isinstance(ast, Name):
if ast.id not in ['True', 'False']:
vars_found.append(ast)
ast = AST.recurse_down_tree(ast, _extract_vars_ast, (vars_found, ))
return ast
def _extract_funcs_ast(ast, funcs_found):
"""
Append ('name', #arg) for each function used in the ast to funcs_found.
"""
if isinstance(ast, Call):
funcs_found.append((AST.ast2str(ast.func), len(ast.args)))
for node in ast.args:
_extract_funcs_ast(node, funcs_found)
ast = AST.recurse_down_tree(ast, _extract_funcs_ast, (funcs_found, ))
return ast
def _make_c_compatible_ast(ast):
if isinstance(ast, BinOp) and isinstance(ast.op, Pow):
ast = Call(func=Name(id='pow', ctx=Load()), args=[ast.left, ast.right], keywords=[])
ast = AST.recurse_down_tree(ast, _make_c_compatible_ast)
elif isinstance(ast, Constant) and isinstance(ast.value, int):
ast.value = float(ast.value)
elif isinstance(ast, Subscript):
# These asts correspond to array[blah] and we shouldn't convert these
# to floats, so we don't recurse down the tree in this case.
pass
# We need to subsitute the C logical operators. Unfortunately, they aren't
# valid python syntax, so we have to cheat a little, using Compare and Name
# nodes abusively. This abuse may not be future-proof... sigh...
elif isinstance(ast, BoolOp) and isinstance(ast.op, And):
nodes = AST.recurse_down_tree(ast.values, _make_c_compatible_ast)
ops = [('&&', node) for node in nodes[1:]]
ops_1 = []
c = []
for k, v in ops:
ops_1.append(k)
c.append(v)
ast = Compare(nodes[0], ops_1, c)
elif isinstance(ast, BoolOp) and isinstance(ast.op, Or):
nodes = AST.recurse_down_tree(ast.values, _make_c_compatible_ast)
ops = [('||', node) for node in nodes[1:]]
ops_1 = []
c = []
for k, v in ops:
ops_1.append(k)
c.append(v)
ast = AST.Compare(nodes[0], ops_1, c)
elif isinstance(ast, UnaryOp) and isinstance(ast.op, Not):
expr = AST.recurse_down_tree(ast.operand, _make_c_compatible_ast)
ast = AST.Name(id='!(%s)' % ast2str(expr))
else:
ast = AST.recurse_down_tree(ast, _make_c_compatible_ast)
return ast
def _sub_subtrees_for_comps(ast, ast_mappings):
if isinstance(ast, Compare) and ast2str(ast) in ast_mappings:
return ast_mappings[ast2str(ast)]
ast = AST.recurse_down_tree(ast, _sub_subtrees_for_comps, (ast_mappings,))
return ast