diff --git a/lark/visitors.py b/lark/visitors.py index 376732e9f..ff6e262f3 100644 --- a/lark/visitors.py +++ b/lark/visitors.py @@ -326,6 +326,8 @@ def transform(self, tree: Tree[_Leaf_T]) -> _Return_T: else: stack.append(x) + if not stack: + return None # type: ignore[return-value] result, = stack # We should have only one tree remaining # There are no guarantees on the type of the value produced by calling a user func for a # child will produce. This means type system can't statically know that the final result is diff --git a/tests/test_trees.py b/tests/test_trees.py index c0cf2b48e..10afc5219 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -415,6 +415,17 @@ def IGNORE_TOKEN(self, token): result = T().transform(copied) self.assertEqual(result, Tree('start', [3, 7])) + def test_transformer_discard_root(self): + """Transformer_NonRecursive should return None (not crash) when the root node is discarded.""" + for base in (Transformer, Transformer_NonRecursive): + class T(base): + def start(self, children): + return Discard + + tree = Tree('start', [Token('A', 'x')]) + result = T().transform(tree) + self.assertIsNone(result) + def test_merge_transformers(self): tree = Tree('start', [ Tree('main', [