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