Package SloppyCell :: Package ExprManip :: Module Substitution
[hide private]

Source Code for Module SloppyCell.ExprManip.Substitution

  1  from compiler.ast import * 
  2  import copy 
  3   
  4  import AST 
  5  from AST import strip_parse, ast2str 
  6  import Simplify 
  7   
8 -def sub_for_comps(expr, mapping):
9 """ 10 For each pair out_name:in_expr in mapping, the returned string has all 11 occurences of the variable out_compe substituted by in_expr. 12 """ 13 if len(mapping) == 0: 14 return expr 15 16 ast = strip_parse(expr) 17 ast_mapping = {} 18 for out_expr, in_expr in mapping.items(): 19 out_ast = strip_parse(out_expr) 20 if not isinstance(out_ast, Compare): 21 raise ValueError, 'Expression %s to substitute for is not a '\ 22 'comparison.' % out_expr 23 ast_mapping[ast2str(out_ast)] = strip_parse(in_expr) 24 25 ast = _sub_subtrees_for_comps(ast, ast_mapping) 26 return ast2str(ast)
27
28 -def _sub_subtrees_for_comps(ast, ast_mappings):
29 if isinstance(ast, Compare) and ast_mappings.has_key(ast2str(ast)): 30 return ast_mappings[ast2str(ast)] 31 ast = AST.recurse_down_tree(ast, _sub_subtrees_for_comps, (ast_mappings,)) 32 return ast
33
34 -def sub_for_var(expr, out_name, in_expr):
35 """ 36 Returns a string with all occurances of the variable out_name substituted by 37 in_expr. 38 39 Perhaps regular expressions could do this more simply... 40 """ 41 return sub_for_vars(expr, {out_name:in_expr})
42
43 -def sub_for_vars(expr, mapping):
44 """ 45 For each pair out_name:in_expr in mapping, the returned string has all 46 occurences of the variable out_name substituted by in_expr. 47 """ 48 if len(mapping) == 0: 49 return expr 50 51 ast = strip_parse(expr) 52 ast_mapping = {} 53 for out_name, in_expr in mapping.items(): 54 out_ast = strip_parse(out_name) 55 if not isinstance(out_ast, Name): 56 raise ValueError, 'Expression %s to substitute for is not a '\ 57 'variable name.' % out_name 58 ast_mapping[str(out_ast.name)] = strip_parse(in_expr) 59 60 ast = _sub_subtrees_for_vars(ast, ast_mapping) 61 return ast2str(ast)
62
63 -def _sub_subtrees_for_vars(ast, ast_mappings):
64 """ 65 For each out_name, in_ast pair in mappings, substitute in_ast for all 66 occurances of the variable named out_name in ast 67 """ 68 if isinstance(ast, Name) and ast_mappings.has_key(ast2str(ast)): 69 return ast_mappings[ast2str(ast)] 70 ast = AST.recurse_down_tree(ast, _sub_subtrees_for_vars, (ast_mappings,)) 71 return ast
72
73 -def sub_for_func(expr, func_name, func_vars, func_expr):
74 """ 75 Return a string with the function func_name substituted for its exploded 76 form. 77 78 func_name: The name of the function. 79 func_vars: A sequence variables used by the function expression 80 func_expr: The expression for the function. 81 For example: 82 If f(x, y, z) = sqrt(z)*x*y-z 83 func_name = 'f' 84 func_vars = ['x', 'y', 'z'] 85 func_expr = 'sqrt(z)*x*y-z' 86 87 As a special case, functions that take a variable number of arguments can 88 use '*' for func_vars. 89 For example: 90 sub_for_func('or_func(or_func(A,D),B,C)', 'or_func', '*', 'x or y') 91 yields '(A or D) or B or C' 92 """ 93 ast = strip_parse(expr) 94 func_name_ast = strip_parse(func_name) 95 if not isinstance(func_name_ast, Name): 96 raise ValueError, 'Function name is not a simple name.' 97 func_name = func_name_ast.name 98 99 func_expr_ast = strip_parse(func_expr) 100 # We can strip_parse the '*', so we special case it here. 101 if func_vars == '*': 102 if not hasattr(func_expr_ast, 'nodes'): 103 raise ValueError("Top-level function in %s does not appear to " 104 "accept variable number of arguments. (It has no " 105 "'nodes' attribute.)" % func_expr) 106 107 func_var_names = '*' 108 else: 109 func_vars_ast = [strip_parse(var) for var in func_vars] 110 for var_ast in func_vars_ast: 111 if not isinstance(var_ast, Name): 112 raise ValueError, 'Function variable is not a simple name.' 113 func_var_names = [getattr(var_ast, 'name') for var_ast in func_vars_ast] 114 115 ast = _sub_for_func_ast(ast, func_name, func_var_names, func_expr_ast) 116 simple = Simplify._simplify_ast(ast) 117 return ast2str(simple)
118
119 -def _sub_for_func_ast(ast, func_name, func_vars, func_expr_ast):
120 """ 121 Return an ast with the function func_name substituted out. 122 """ 123 if isinstance(ast, CallFunc) and ast2str(ast.node) == func_name\ 124 and func_vars == '*': 125 working_ast = copy.deepcopy(func_expr_ast) 126 new_args = [_sub_for_func_ast(arg_ast, func_name, func_vars, 127 func_expr_ast) for arg_ast in ast.args] 128 # This subs out the arguments of the original function. 129 working_ast.nodes = new_args 130 return working_ast 131 if isinstance(ast, CallFunc) and ast2str(ast.node) == func_name\ 132 and len(ast.args) == len(func_vars): 133 # If our ast is the function we're looking for, we take the ast 134 # for the function expression, substitute for its arguments, and 135 # return 136 working_ast = copy.deepcopy(func_expr_ast) 137 mapping = {} 138 for var_name, arg_ast in zip(func_vars, ast.args): 139 subbed_arg_ast = _sub_for_func_ast(arg_ast, func_name, func_vars, 140 func_expr_ast) 141 mapping[var_name] = subbed_arg_ast 142 _sub_subtrees_for_vars(working_ast, mapping) 143 return working_ast 144 ast = AST.recurse_down_tree(ast, _sub_for_func_ast, 145 (func_name, func_vars, func_expr_ast,)) 146 return ast
147
148 -def make_c_compatible(expr):
149 """ 150 Convert a python math string into one compatible with C. 151 152 Substitute all python-style x**n exponents with pow(x, n). 153 Replace all integer constants with float values to avoid integer 154 casting problems (e.g. '1' -> '1.0'). 155 Replace 'and', 'or', and 'not' with C's '&&', '||', and '!'. This may be 156 fragile if the parsing library changes in newer python versions. 157 """ 158 ast = strip_parse(expr) 159 ast = _make_c_compatible_ast(ast) 160 return ast2str(ast)
161
162 -def _make_c_compatible_ast(ast):
163 if isinstance(ast, Power): 164 ast = CallFunc(Name('pow'), [ast.left, ast.right], None, None) 165 ast = AST.recurse_down_tree(ast, _make_c_compatible_ast) 166 elif isinstance(ast, Const) and isinstance(ast.value, int): 167 ast.value = float(ast.value) 168 elif isinstance(ast, Subscript): 169 # These asts correspond to array[blah] and we shouldn't convert these 170 # to floats, so we don't recurse down the tree in this case. 171 pass 172 # We need to subsitute the C logical operators. Unfortunately, they aren't 173 # valid python syntax, so we have to cheat a little, using Compare and Name 174 # nodes abusively. This abuse may not be future-proof... sigh... 175 elif isinstance(ast, And): 176 nodes = AST.recurse_down_tree(ast.nodes, _make_c_compatible_ast) 177 ops = [('&&', node) for node in nodes[1:]] 178 ast = AST.Compare(nodes[0], ops) 179 elif isinstance(ast, Or): 180 nodes = AST.recurse_down_tree(ast.nodes, _make_c_compatible_ast) 181 ops = [('||', node) for node in nodes[1:]] 182 ast = AST.Compare(nodes[0], ops) 183 elif isinstance(ast, Not): 184 expr = AST.recurse_down_tree(ast.expr, _make_c_compatible_ast) 185 ast = AST.Name('!(%s)' % ast2str(expr)) 186 else: 187 ast = AST.recurse_down_tree(ast, _make_c_compatible_ast) 188 return ast
189