[Python-checkins] r62470 - in sandbox/trunk/2to3/lib2to3: refactor.py tests/test_fixers.py

[david.wolever]

Author: david.wolever
Date: Thu Apr 24 02:11:07 2008
New Revision: 62470

Log:
Fixed up and applied the patch for #2431 -- speeding up 2to3 with a lookup table.



Modified:
   sandbox/trunk/2to3/lib2to3/refactor.py
   sandbox/trunk/2to3/lib2to3/tests/test_fixers.py

Modified: sandbox/trunk/2to3/lib2to3/refactor.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/refactor.py	(original)
+++ sandbox/trunk/2to3/lib2to3/refactor.py	Thu Apr 24 02:11:07 2008
@@ -18,6 +18,8 @@
 import difflib
 import optparse
 import logging
+from collections import defaultdict
+from itertools import chain
 
 # Local imports
 from .pgen2 import driver
@@ -96,6 +98,43 @@
     fix_names.sort()
     return fix_names
 
+def get_head_types(pat):
+    """ Accepts a pytree Pattern Node and returns a set
+        of the pattern types which will match first. """
+
+    if isinstance(pat, (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 isinstance(pat, pytree.NegatedPattern):
+        if pat.content:
+            return get_head_types(pat.content)
+        return set([None]) # Negated Patterns don't have a type
+
+    if isinstance(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))
+
+def get_headnode_dict(fixer_list):
+    """ Accepts a list of fixers and returns a dictionary
+        of head node type --> fixer list.  """
+    head_nodes = defaultdict(list)
+    for fixer in fixer_list:
+        if not fixer.pattern:
+            head_nodes[None].append(fixer)
+            continue
+        for t in get_head_types(fixer.pattern):
+            head_nodes[t].append(fixer)
+    return head_nodes
+
 
 class RefactoringTool(object):
 
@@ -114,6 +153,10 @@
                                     convert=pytree.convert,
                                     logger=self.logger)
         self.pre_order, self.post_order = self.get_fixers()
+
+        self.pre_order = get_headnode_dict(self.pre_order)
+        self.post_order = get_headnode_dict(self.post_order)
+
         self.files = []  # List of files that were or should be modified
 
     def get_fixers(self):
@@ -286,7 +329,11 @@
         Returns:
             True if the tree was modified, False otherwise.
         """
-        all_fixers = self.pre_order + self.post_order
+        # Two calls to chain are required because pre_order.values()
+        #   will be a list of lists of fixers:
+        #   [[<fixer ...>, <fixer ...>], [<fixer ...>]]
+        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 +359,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)

Modified: sandbox/trunk/2to3/lib2to3/tests/test_fixers.py
==============================================================================
--- sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	(original)
+++ sandbox/trunk/2to3/lib2to3/tests/test_fixers.py	Thu Apr 24 02:11:07 2008
@@ -33,8 +33,10 @@
         self.fixer_log = []
         self.filename = "<string>"
 
-        for order in (self.refactor.pre_order, self.refactor.post_order):
-            for fixer in order:
+        from itertools import chain
+        for order in (self.refactor.pre_order.values(),\
+                      self.refactor.post_order.values()):
+            for fixer in chain(*order):
                 fixer.log = self.fixer_log
 
     def _check(self, before, after):