1 #!/usr/bin/env python 2 3 import compiler 4 5 def process_nodes(root_node, prefix): 6 7 """ 8 Under the 'root_node', process all suitable expression nodes which employ 9 Name nodes using the given 'prefix'. 10 """ 11 12 for node in root_node.getChildNodes(): 13 # Identify suitable names and add such nodes 14 if isinstance(node, compiler.ast.CallFunc): 15 process_callfunc(node, prefix) 16 elif isinstance(node, compiler.ast.Getattr): 17 process_getattr(node, prefix, root_node) 18 else: 19 process_nodes(node, prefix) 20 21 def process_callfunc(node, prefix): 22 # Check the target. 23 target = node.node 24 if isinstance(target, compiler.ast.Getattr): 25 process_getattr(target, prefix, node) 26 else: 27 process_nodes(node, prefix) 28 29 def process_getattr(node, prefix, parent): 30 # Check the target. 31 target = node.expr 32 if isinstance(target, compiler.ast.Name) and target.name.startswith(prefix): 33 34 # Replace CallFunc plus Getattr occurrences: 35 # node.attr(args) -> Node_attr(node, args) 36 37 if isinstance(parent, compiler.ast.CallFunc): 38 parent.node = compiler.ast.Name("Node_%s" % node.attrname) 39 parent.args.insert(0, compiler.ast.Name(target.name)) 40 41 # Replace plain Getattr nodes: 42 # node.attr -> Node_attr(node) 43 # NOTE: Nasty but necessary rewiring of the parent node required. 44 45 else: 46 replacement = compiler.ast.CallFunc( 47 compiler.ast.Name("Node_%s" % node.attrname), 48 [node.expr] 49 ) 50 for key, value in parent.__dict__.items(): 51 # Detect lists. 52 try: 53 if node in value: 54 index = value.index(node) 55 value[index] = replacement 56 except TypeError: 57 if value is node: 58 parent.__dict__[key] = replacement 59 else: 60 process_nodes(node, prefix) 61 62 def process_file(filename, prefix="x2_"): 63 64 """ 65 Process the module given by the specified 'filename'. The optional special 66 'prefix' marks those variables to be processed. 67 """ 68 69 # Open the module as an AST. 70 71 module = compiler.parseFile(filename) 72 73 # Find references to special variables. 74 75 process_nodes(module, prefix) 76 return module 77 78 # vim: tabstop=4 expandtab shiftwidth=4