diff --git a/lark/visitors.py b/lark/visitors.py index 03dc7404..ae7930c0 100644 --- a/lark/visitors.py +++ b/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) diff --git a/tests/test_trees.py b/tests/test_trees.py index 36be73fb..7d2cc7f8 100644 --- a/tests/test_trees.py +++ b/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]))