Fixed Discard, broken in 0.6.6 (Issue #337)

lark/visitors.py

@@ -40,7 +40,7 @@ def _call_userfunc(self, tree, new_children=None):
                     return f(tree)
                 else:
                     return f(children)
-            except GrammarError:
+            except (GrammarError, Discard):
                 raise
             except Exception as e:
                 raise VisitError(tree, e)

tests/test_trees.py

@@ -6,7 +6,7 @@
 import pickle
 
 from lark.tree import Tree
-from lark.visitors import Transformer, Interpreter, visit_children_decor, v_args
+from lark.visitors import Transformer, Interpreter, visit_children_decor, v_args, Discard
 
 
 class TestTrees(TestCase):
@@ -144,6 +144,26 @@ def add(self, values):
         res = T().transform(t)
         self.assertEqual(res, 2.9)
 
+    def test_discard(self):
+        class MyTransformer(Transformer):
+            def a(self, args):
+                return 1 # some code here
+
+            def b(cls, args):
+                raise Discard()
+
+        t = Tree('root', [
+            Tree('b', []),
+            Tree('a', []),
+            Tree('b', []),
+            Tree('c', []),
+            Tree('b', []),
+        ])
+        t2 = Tree('root', [1, Tree('c', [])])
+
+        x = MyTransformer().transform( t )
+        self.assertEqual(x, t2)
+
 
 if __name__ == '__main__':
     unittest.main()