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()