Refactored v_args & visitors to a better, more agile implementation

lark/parse_tree_builder.py

@@ -3,6 +3,7 @@ from .lexer import Token
 from .tree import Tree
 from .visitors import InlineTransformer # XXX Deprecated
 from .visitors import Transformer_InPlace
+from . import visitors
 
 ###{standalone
 from functools import partial, wraps
@@ -202,6 +203,15 @@ def inplace_transformer(func):
         return func(tree)
     return f
 
+def apply_visit_wrapper(func, name, wrapper):
+    if wrapper is visitors._vargs_meta or wrapper is visitors._vargs_meta_inline:
+        raise NotImplementedError("Meta args not supported for internal transformer")
+    @wraps(func)
+    def f(children):
+        return wrapper(func, name, children, None)
+    return f
+
+
 class ParseTreeBuilder:
     def __init__(self, rules, tree_class, propagate_positions=False, keep_all_tokens=False, ambiguous=False, maybe_placeholders=False):
         self.tree_class = tree_class
@@ -236,12 +246,15 @@ class ParseTreeBuilder:
             user_callback_name = rule.alias or rule.origin.name
             try:
                 f = getattr(transformer, user_callback_name)
-                assert not getattr(f, 'meta', False), "Meta args not supported for internal transformer"
                 # XXX InlineTransformer is deprecated!
-                if getattr(f, 'inline', False) or isinstance(transformer, InlineTransformer):
-                    f = ptb_inline_args(f)
-                elif hasattr(f, 'whole_tree') or isinstance(transformer, Transformer_InPlace):
-                    f = inplace_transformer(f)
+                wrapper = getattr(f, 'visit_wrapper', None)
+                if wrapper is not None:
+                    f = apply_visit_wrapper(f, user_callback_name, wrapper)
+                else:
+                    if isinstance(transformer, InlineTransformer):
+                        f = ptb_inline_args(f)
+                    elif isinstance(transformer, Transformer_InPlace):
+                        f = inplace_transformer(f)
             except AttributeError:
                 f = partial(self.tree_class, user_callback_name)
 

lark/visitors.py

@@ -35,14 +35,9 @@ class Transformer:
             return self.__default__(tree.data, children, tree.meta)
         else:
             try:
-                if getattr(f, 'meta', False):
-                    return f(children, tree.meta)
-                elif getattr(f, 'inline', False):
-                    return f(*children)
-                elif getattr(f, 'whole_tree', False):
-                    if new_children is not None:
-                        tree.children = new_children
-                    return f(tree)
+                wrapper = getattr(f, 'visit_wrapper', None)
+                if wrapper is not None:
+                    return f.visit_wrapper(f, tree.data, children, tree.meta)
                 else:
                     return f(children)
             except (GrammarError, Discard):
@@ -282,8 +277,7 @@ def inline_args(obj):   # XXX Deprecated
 
 
 
-def _visitor_args_func_dec(func, inline=False, meta=False, whole_tree=False, static=False):
-    assert [whole_tree, meta, inline].count(True) <= 1
+def _visitor_args_func_dec(func, visit_wrapper=None, static=False):
     def create_decorator(_f, with_self):
         if with_self:
             def f(self, *args, **kwargs):
@@ -298,17 +292,42 @@ def _visitor_args_func_dec(func, inline=False, meta=False, whole_tree=False, sta
     else:
         f = smart_decorator(func, create_decorator)
     f.vargs_applied = True
-    f.inline = inline
-    f.meta = meta
-    f.whole_tree = whole_tree
+    f.visit_wrapper = visit_wrapper
     return f
 
-def v_args(inline=False, meta=False, tree=False):
+
+def _vargs_inline(f, data, children, meta):
+    return f(*children)
+def _vargs_meta_inline(f, data, children, meta):
+    return f(meta, *children)
+def _vargs_meta(f, data, children, meta):
+    return f(children, meta)   # TODO swap these for consistency? Backwards incompatible!
+def _vargs_tree(f, data, children, meta):
+    return f(Tree(data, children, meta))
+
+def v_args(inline=False, meta=False, tree=False, wrapper=None):
     "A convenience decorator factory, for modifying the behavior of user-supplied visitor methods"
-    if [tree, meta, inline].count(True) > 1:
-        raise ValueError("Visitor functions can either accept tree, or meta, or be inlined. These cannot be combined.")
+    if tree and (meta or inline):
+        raise ValueError("Visitor functions cannot combine 'tree' with 'meta' or 'inline'.")
+
+    func = None
+    if meta:
+        if inline:
+            func = _vargs_meta_inline
+        else:
+            func = _vargs_meta
+    elif inline:
+        func = _vargs_inline
+    elif tree:
+        func = _vargs_tree
+
+    if wrapper is not None:
+        if func is not None:
+            raise ValueError("Cannot use 'wrapper' along with 'tree', 'meta' or 'inline'.")
+        func = wrapper
+
     def _visitor_args_dec(obj):
-        return _apply_decorator(obj, _visitor_args_func_dec, inline=inline, meta=meta, whole_tree=tree)
+        return _apply_decorator(obj, _visitor_args_func_dec, visit_wrapper=func)
     return _visitor_args_dec
 
 

tests/test_parser.py

@@ -5,6 +5,7 @@ import unittest
 import logging
 import os
 import sys
+from copy import deepcopy
 try:
     from cStringIO import StringIO as cStringIO
 except ImportError:
@@ -117,6 +118,61 @@ class TestParsers(unittest.TestCase):
         r = p.parse("x")
         self.assertEqual( r.children, ["X!"] )
 
+    def test_vargs_meta(self):
+
+        @v_args(meta=True)
+        class T1(Transformer):
+            def a(self, children, meta):
+                assert not children
+                return meta.line
+
+            def start(self, children, meta):
+                return children
+
+        @v_args(meta=True, inline=True)
+        class T2(Transformer):
+            def a(self, meta):
+                return meta.line
+
+            def start(self, meta, *res):
+                return list(res)
+
+        for T in (T1, T2):
+            for internal in [False, True]:
+                try:
+                    g = Lark(r"""start: a+
+                                a : "x" _NL?
+                                _NL: /\n/+
+                            """, parser='lalr', transformer=T() if internal else None)
+                except NotImplementedError:
+                    assert internal
+                    continue
+
+                res = g.parse("xx\nx\nxxx\n\n\nxx")
+                assert not internal
+                res = T().transform(res)
+
+                self.assertEqual(res, [1, 1, 2, 3, 3, 3, 6, 6])
+
+    def test_vargs_tree(self):
+        tree = Lark('''
+            start: a a a
+            !a: "A"
+        ''').parse('AAA')
+        tree_copy = deepcopy(tree)
+
+        @v_args(tree=True)
+        class T(Transformer):
+            def a(self, tree):
+                return 1
+            def start(self, tree):
+                return tree.children
+
+        res = T().transform(tree)
+        self.assertEqual(res, [1, 1, 1])
+        self.assertEqual(tree, tree_copy)
+
+
 
     def test_embedded_transformer(self):
         class T(Transformer):
@@ -188,7 +244,7 @@ class TestParsers(unittest.TestCase):
         @v_args(tree=True)
         class T2(Transformer):
             def a(self, tree):
-                assert isinstance(tree, Tree)
+                assert isinstance(tree, Tree), tree
                 tree.children.append("tested")
                 return tree