diff --git a/lark/visitors.py b/lark/visitors.py index f89d942..26ca242 100644 --- a/lark/visitors.py +++ b/lark/visitors.py @@ -40,7 +40,7 @@ class Transformer: return f(tree) else: return f(children) - except GrammarError: + except (GrammarError, Discard): raise except Exception as e: raise VisitError(tree, e) diff --git a/tests/test_trees.py b/tests/test_trees.py index 72cfbf1..e5008ef 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -6,7 +6,7 @@ import copy 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 @@ class TestTrees(TestCase): 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()