m-labs · May 28, 2015 · May 28, 2015
Showing with 71 additions and 14 deletions.

+46 −12 pythonparser/algorithm.py

+5 −2 pythonparser/ast.py

+20 −0 pythonparser/test/test_algorithm.py
diff --git a/pythonparser/algorithm.py b/pythonparser/algorithm.py
@@ -8,36 +8,70 @@
 
 class Visitor:
     """
-    A node visitor base class that does a pre-order traversal
+    A node visitor base class that does a traversal
     of the abstract syntax tree.
 
-    This class is meant to be subclassed, with the subclass adding visitor
-    methods.
+    This class is meant to be subclassed, with the subclass adding
+    visitor methods. The visitor method should call ``self.generic_visit(node)``
+    to continue the traversal; this allows to perform arbitrary
+    actions both before and after traversing the children of a node.
 
-    The visitor functions for the nodes are ``'visit_'`` +
+    The visitor methods for the nodes are ``'visit_'`` +
     class name of the node.  So a `Try` node visit function would
     be `visit_Try`.
     """
 
     def generic_visit(self, node):
         """Called if no explicit visitor function exists for a node."""
-        pass
+        for field_name in node._fields:
+            field_val = getattr(node, field_name)
+            if isinstance(field_val, list):
+                for field_val_elt in field_val:
+                    self.visit(field_val_elt)
+            elif isinstance(field_val, ast.AST):
+                self.visit(field_val)
 
     def visit(self, node):
         """Visit a node."""
         visit_attr = 'visit_' + type(node).__name__
         if hasattr(self, visit_attr):
-            getattr(self, visit_attr)(node)
+            return getattr(self, visit_attr)(node)
         else:
-            self.generic_visit(node)
+            return self.generic_visit(node)
 
+class Transformer(Visitor):
+    """
+    A node transformer base class that does a post-order traversal
+    of the abstract syntax tree while allowing to replace or remove
+    the nodes being traversed.
+
+    The return value of the visitor methods is used to replace or remove
+    the old node.  If the return value of the visitor method is ``None``,
+    the node will be removed from its location, otherwise it is replaced
+    with the return value.  The return value may be the original node
+    in which case no replacement takes place.
+
+    This class is meant to be subclassed, with the subclass adding
+    visitor methods. The visitor method should call ``self.generic_visit(node)``
+    to continue the traversal; this allows to perform arbitrary
+    actions both before and after traversing the children of a node.
+
+    The visitor methods for the nodes are ``'visit_'`` +
+    class name of the node.  So a `Try` node visit function would
+    be `visit_Try`.
+    """
+
+    def generic_visit(self, node):
+        """Called if no explicit visitor function exists for a node."""
         for field_name in node._fields:
             field_val = getattr(node, field_name)
-            if isinstance(field_val, ast.AST):
-                self.visit(field_val)
-            elif isinstance(field_val, list):
-                for field_val_elt in field_val:
-                    self.visit(field_val_elt)
+            if isinstance(field_val, list):
+                setattr(node, field_name,
+                        list(filter(lambda x: x is not None, map(self.visit, field_val))))
+            elif isinstance(field_val, ast.AST):
+                setattr(node, field_name, self.visit(field_val))
+
+        return node
 
 def compare(left, right, compare_locs=False):
     """

diff --git a/pythonparser/ast.py b/pythonparser/ast.py
@@ -40,6 +40,9 @@ class commonloc(object):
 
     _locs = ('loc',)
 
+    def _reprfields(self):
+        return self._fields + self._locs
+
     def __repr__(self):
         def value(name):
             try:
@@ -51,7 +54,7 @@ def value(name):
             except:
                 return "(!!!MISSING!!!)"
         fields = ', '.join(map(lambda name: "%s=%s" % (name, value(name)),
-                           self._fields + self._locs))
+                           self._reprfields()))
         return "%s(%s)" % (self.__class__.__name__, fields)
 
 class keywordloc(commonloc):
@@ -73,7 +76,7 @@ class beginendloc(commonloc):
 
 # AST nodes
 
-class AST:
+class AST(object):
     """
     An ancestor of all nodes.
 

diff --git a/pythonparser/test/test_algorithm.py b/pythonparser/test/test_algorithm.py
@@ -11,9 +11,11 @@ def __init__(self):
 
             def visit_Num(self, node):
                 self.num_count += 1
+                algorithm.Visitor.generic_visit(self, node)
 
             def generic_visit(self, node):
                 self.other_count += 1
+                algorithm.Visitor.generic_visit(self, node)
 
         visitor = FooVisitor()
         visitor.visit(parse("[1,2,x,y]\n"))
@@ -27,3 +29,21 @@ def test_compare(self):
         self.assertTrue( algorithm.compare(parse("1 + 2\n"), parse("1 +  2\n")))
         self.assertFalse(algorithm.compare(parse("1 + 2\n"), parse("1 +  2\n"),
                                            compare_locs=True))
+
+    def test_transform(self):
+        class FooTransformer(algorithm.Transformer):
+            def visit_Num(self, node):
+                if node.n == 42:
+                    return None
+                return node
+
+        transformer = FooTransformer()
+
+        self.assertTrue(algorithm.compare(parse("return 10\n"),
+                        transformer.visit(parse("return 10\n"))))
+        self.assertTrue(algorithm.compare(parse("return\n"),
+                        transformer.visit(parse("return 42\n"))))
+        self.assertTrue(algorithm.compare(parse("[1,2,3]\n"),
+                        transformer.visit(parse("[1,2,3]\n"))))
+        self.assertTrue(algorithm.compare(parse("[1,3]\n"),
+                        transformer.visit(parse("[1,42,3]\n"))))