diff --git a/lark/visitors.py b/lark/visitors.py index ae9d128c..2d30c53e 100644 --- a/lark/visitors.py +++ b/lark/visitors.py @@ -158,7 +158,11 @@ def _transform_tree(self, tree): def transform(self, tree: Tree[_Leaf_T]) -> _Return_T: "Transform the given tree, and return the final result" - return self._transform_tree(tree) + res = list(self._transform_children([tree])) + if not res: + return None # type: ignore[return-value] + assert len(res) == 1 + return res[0] def __mul__( self: 'Transformer[_Leaf_T, Tree[_Leaf_U]]', diff --git a/tests/test_trees.py b/tests/test_trees.py index dd95f6b8..1f69869e 100644 --- a/tests/test_trees.py +++ b/tests/test_trees.py @@ -447,5 +447,20 @@ class T4(Transformer): with self.assertRaises(AttributeError): merge_transformers(T1(), module=T3()) + def test_transform_token(self): + class MyTransformer(Transformer): + def INT(self, value): + return int(value) + + t = Token('INT', '123') + assert MyTransformer().transform(t) == 123 + + class MyTransformer(Transformer): + def INT(self, value): + return Discard + + assert MyTransformer().transform(t) is None + + if __name__ == '__main__': unittest.main()