Index: lib2to3/refactor.py =================================================================== --- lib2to3/refactor.py (revision 61673) +++ lib2to3/refactor.py (working copy) @@ -114,6 +114,50 @@ convert=pytree.convert, logger=self.logger) self.pre_order, self.post_order = self.get_fixers() + # XXX: REFACTOR + from itertools import chain + def get_head_types(pat): + if type(pat) in (pytree.NodePattern, pytree.LeafPattern): + # NodePatters must either have no type and no content + # or a type and content -- so they don't get any farther + # Always return leafs + return set([pat.type]) + + if type(pat) == pytree.NegatedPattern: + if pat.content: + return get_head_types(pat.content) + return set([None]) # Negated Patterns don't have a type + + if type(pat) == pytree.WildcardPattern: + # Recurse on each node in content + r = set() + for p in pat.content: + for x in p: + r.update(get_head_types(x)) + return r + + raise Exception("Oh no! I don't understand pattern %s" %(pat)) + + from collections import defaultdict + head_nodes = defaultdict(list) + for fixer in self.pre_order: + if not fixer.pattern: + head_nodes[None].append(fixer) + continue + for t in get_head_types(fixer.pattern): + head_nodes[t].append(fixer) + self.pre_order = head_nodes + + head_nodes = defaultdict(list) + for fixer in self.post_order: + if not fixer.pattern: + head_nodes[None].append(fixer) + continue + for t in get_head_types(fixer.pattern): + head_nodes[t].append(fixer) + self.post_order = head_nodes + # XXX: REFACTOR + self.files = [] # List of files that were or should be modified def get_fixers(self): @@ -286,7 +330,8 @@ Returns: True if the tree was modified, False otherwise. """ - all_fixers = self.pre_order + self.post_order + from itertools import chain + all_fixers = chain(chain(*self.pre_order.values()), chain(*self.post_order.values())) for fixer in all_fixers: fixer.start_tree(tree, name) @@ -312,7 +357,7 @@ if not fixers: return for node in traversal: - for fixer in fixers: + for fixer in fixers[node.type] + fixers[None]: results = fixer.match(node) if results: new = fixer.transform(node, results)