From 2b3fabd01cec3f6803ec7d255bf0436476c2b812 Mon Sep 17 00:00:00 2001 From: Erez Shinan <erezshin+git@gmail.com> Date: Sat, 9 Mar 2019 13:04:17 +0200 Subject: [PATCH] Fixed Discard, broken in 0.6.6 (Issue #337) --- lark/visitors.py | 2 +- tests/test_trees.py | 22 +++++++++++++++++++++- 2 files changed, 22 insertions(+), 2 deletions(-) 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()