Allow transformers to transform tokens (Issue #389)

lark/visitors.py

@@ -3,6 +3,7 @@ from functools import wraps
 from .utils import smart_decorator
 from .tree import Tree
 from .exceptions import VisitError, GrammarError
+from .lexer import Token
 
 ###{standalone
 from inspect import getmembers, getmro
@@ -21,6 +22,10 @@ class Transformer:
     Can be used to implement map or reduce.
     """
 
+    __visit_tokens__ = False   # For backwards compatibility
+    def __init__(self,  visit_tokens=False):
+        self.__visit_tokens__ = visit_tokens
+
     def _call_userfunc(self, tree, new_children=None):
         # Assumes tree is already transformed
         children = new_children if new_children is not None else tree.children
@@ -45,10 +50,29 @@ class Transformer:
             except Exception as e:
                 raise VisitError(tree, e)
 
+    def _call_userfunc_token(self, token):
+        try:
+            f = getattr(self, token.type)
+        except AttributeError:
+            return self.__default_token__(token)
+        else:
+            try:
+                return f(token)
+            except (GrammarError, Discard):
+                raise
+            except Exception as e:
+                raise VisitError(token, e)
+
+
     def _transform_children(self, children):
         for c in children:
             try:
-                yield self._transform_tree(c) if isinstance(c, Tree) else c
+                if isinstance(c, Tree):
+                    yield self._transform_tree(c)
+                elif self.__visit_tokens__ and isinstance(c, Token):
+                    yield self._call_userfunc_token(c)
+                else:
+                    yield c
             except Discard:
                 pass
 
@@ -66,6 +90,11 @@ class Transformer:
         "Default operation on tree (for override)"
         return Tree(data, children, meta)
 
+    def __default_token__(self, token):
+        "Default operation on token (for override)"
+        return token
+
+
     @classmethod
     def _apply_decorator(cls, decorator, **kwargs):
         mro = getmro(cls)

tests/test_parser.py

@@ -94,6 +94,24 @@ class TestParsers(unittest.TestCase):
         r = g.parse('xx')
         self.assertEqual( r.children[0].data, "c" )
 
+    def test_visit_tokens(self):
+        class T(Transformer):
+            def a(self, children):
+                return children[0] + "!"
+            def A(self, tok):
+                return tok.upper()
+
+        # Test regular
+        g = Lark("""start: a
+                    a : A
+                    A: "x"
+                 """, parser='lalr')
+        r = T().transform(g.parse("x"))
+        self.assertEqual( r.children, ["x!"] )
+        r = T(True).transform(g.parse("x"))
+        self.assertEqual( r.children, ["X!"] )
+
+
     def test_embedded_transformer(self):
         class T(Transformer):
             def a(self, children):