Merge pull request #1013 from adbancroft/nonrecursive_discard

lark/visitors.py

@@ -277,9 +277,15 @@ def transform(self, tree):
                     del stack[-size:]
                 else:
                     args = []
-                stack.append(self._call_userfunc(x, args))
+                try:
+                    stack.append(self._call_userfunc(x, args))
+                except Discard:
+                    pass
             elif self.__visit_tokens__ and isinstance(x, Token):
-                stack.append(self._call_userfunc_token(x))
+                try:
+                    stack.append(self._call_userfunc_token(x))
+                except Discard:
+                    pass
             else:
                 stack.append(x)
 

tests/test_trees.py

@@ -375,7 +375,11 @@ def b(cls, args):
         self.assertEqual(x, t2)
 
     def test_transformer_variants(self):
-        tree = Tree('start', [Tree('add', [Token('N', '1'), Token('N', '2')]), Tree('add', [Token('N', '3'), Token('N', '4')])])
+        tree = Tree('start', [
+            Tree('add', [Token('N', '1'), Token('N', '2'), Token('IGNORE_TOKEN', '4')]),
+            Tree('add', [Token('N', '3'), Token('N', '4')]),
+            Tree('ignore_tree', [Token('DO', 'NOT PANIC')]),
+            ])
         for base in (Transformer, Transformer_InPlace, Transformer_NonRecursive, Transformer_InPlaceRecursive):
             class T(base):
                 def add(self, children):
@@ -384,6 +388,12 @@ def add(self, children):
                 def N(self, token):
                     return int(token)
 
+                def ignore_tree(self, children):
+                    raise Discard
+
+                def IGNORE_TOKEN(self, token):
+                    raise Discard
+
             copied = copy.deepcopy(tree)
             result = T().transform(copied)
             self.assertEqual(result, Tree('start', [3, 7]))