Source Code for Module SloppyCell.ExprManip.Extraction

 1  from compiler.ast import * 
 2  import sets 
 3   
 4  import AST 
 5   
 6  extract_vars_cache = {} 
 7   
 8 -def extract_comps(expr): 
 9      """ 
10      Extract all comparisons from the expression. 
11      """ 
12      comps_found = [] 
13      _extract_comps_ast(AST.strip_parse(expr), comps_found) 
14      comps_found = [AST.ast2str(ast) for ast in comps_found] 
15      return sets.Set(comps_found) 
16   
17 -def _extract_comps_ast(ast, comps_found): 
18      if isinstance(ast, Compare): 
19          comps_found.append(ast) 
20          _extract_comps_ast(ast.expr, comps_found) 
21          for op, elem in ast.ops: 
22              _extract_comps_ast(elem, comps_found) 
23      elif isinstance(ast, list) or isinstance(ast, tuple): 
24          for elem in ast: 
25              _extract_comps_ast(elem, comps_found) 
26      elif AST._node_attrs.has_key(ast.__class__): 
27          for attr_name in AST._node_attrs[ast.__class__]: 
28              attr = getattr(ast, attr_name) 
29              _extract_comps_ast(attr, comps_found) 
30   
31 -def extract_vars(expr): 
32      """ 
33      Return a Set of the variables used in an expression. 
34      """ 
35      try: 
36          return extract_vars_cache[expr] 
37      except KeyError: 
38          vars_found = [] 
39          _extract_vars_ast(AST.strip_parse(expr), vars_found) 
40          vars_found = [AST.ast2str(ast) for ast in vars_found] 
41          result = sets.Set(vars_found) 
42          extract_vars_cache[expr] = result 
43          return result 
44   
45 -def _extract_vars_ast(ast, vars_found): 
46      """ 
47      Appends the asts of the variables used in ast to vars_found. 
48      """ 
49      if isinstance(ast, Name): 
50          if ast.name not in ['True', 'False']: 
51              vars_found.append(ast) 
52      ast = AST.recurse_down_tree(ast, _extract_vars_ast, (vars_found,)) 
53      return ast 
54   
55 -def extract_funcs(expr): 
56      """ 
57      Return a Set of the functions used in an expression. 
58   
59      The elements of the Set are ('function_name', #arguments). 
60      """ 
61      funcs_found = [] 
62      _extract_funcs_ast(AST.strip_parse(expr), funcs_found) 
63      return sets.Set(funcs_found) 
64   
65 -def _extract_funcs_ast(ast, funcs_found): 
66      """ 
67      Append ('name', #arg) for each function used in the ast to funcs_found. 
68      """ 
69      if isinstance(ast, CallFunc): 
70          funcs_found.append((AST.ast2str(ast.node), len(ast.args))) 
71          for node in ast.args: 
72              _extract_funcs_ast(node, funcs_found) 
73      ast = AST.recurse_down_tree(ast, _extract_funcs_ast, (funcs_found,)) 
74      return ast 
75