From a55b7155b51418444f856f55c31abee7a688380f Mon Sep 17 00:00:00 2001
From: Erez Sh <erezshin@gmail.com>
Date: Sat, 8 Feb 2020 05:37:45 +0200
Subject: [PATCH] Added support for v_args in Interpreter (Issue #520)

---
 lark/visitors.py | 57 ++++++++++++++++++++++++++++--------------------
 1 file changed, 33 insertions(+), 24 deletions(-)

diff --git a/lark/visitors.py b/lark/visitors.py
index da6b1d5..30a2a65 100644
--- a/lark/visitors.py
+++ b/lark/visitors.py
@@ -13,7 +13,31 @@ class Discard(Exception):
 
 # Transformers
 
-class Transformer:
+class _Decoratable:
+    @classmethod
+    def _apply_decorator(cls, decorator, **kwargs):
+        mro = getmro(cls)
+        assert mro[0] is cls
+        libmembers = {name for _cls in mro[1:] for name, _ in getmembers(_cls)}
+        for name, value in getmembers(cls):
+
+            # Make sure the function isn't inherited (unless it's overwritten)
+            if name.startswith('_') or (name in libmembers and name not in cls.__dict__):
+                continue
+            if not callable(cls.__dict__[name]):
+                continue
+
+            # Skip if v_args already applied (at the function level)
+            if hasattr(cls.__dict__[name], 'vargs_applied'):
+                continue
+
+            static = isinstance(cls.__dict__[name], (staticmethod, classmethod))
+            setattr(cls, name, decorator(value, static=static, **kwargs))
+        return cls
+
+
+
+class Transformer(_Decoratable):
     """Visits the tree recursively, starting with the leaves and finally the root (bottom-up)
 
     Calls its methods (provided by user via inheritance) according to tree.data
@@ -90,27 +114,6 @@ class Transformer:
         return token
 
 
-    @classmethod
-    def _apply_decorator(cls, decorator, **kwargs):
-        mro = getmro(cls)
-        assert mro[0] is cls
-        libmembers = {name for _cls in mro[1:] for name, _ in getmembers(_cls)}
-        for name, value in getmembers(cls):
-
-            # Make sure the function isn't inherited (unless it's overwritten)
-            if name.startswith('_') or (name in libmembers and name not in cls.__dict__):
-                continue
-            if not callable(cls.__dict__[name]):
-                continue
-
-            # Skip if v_args already applied (at the function level)
-            if hasattr(cls.__dict__[name], 'vargs_applied'):
-                continue
-
-            static = isinstance(cls.__dict__[name], (staticmethod, classmethod))
-            setattr(cls, name, decorator(value, static=static, **kwargs))
-        return cls
-
 
 class InlineTransformer(Transformer):   # XXX Deprecated
     def _call_userfunc(self, tree, new_children=None):
@@ -221,7 +224,7 @@ def visit_children_decor(func):
     return inner
 
 
-class Interpreter:
+class Interpreter(_Decoratable):
     """Top-down visitor, recursive
 
     Visits the tree, starting with the root and finally the leaves (top-down)
@@ -230,8 +233,14 @@ class Interpreter:
     Unlike Transformer and Visitor, the Interpreter doesn't automatically visit its sub-branches.
     The user has to explicitly call visit_children, or use the @visit_children_decor
     """
+
     def visit(self, tree):
-        return getattr(self, tree.data)(tree)
+        f = getattr(self, tree.data)
+        wrapper = getattr(f, 'visit_wrapper', None)
+        if wrapper is not None:
+            return f.visit_wrapper(f, tree.data, tree.children, tree.meta)
+        else:
+            return f(tree)
 
     def visit_children(self, tree):
         return [self.visit(child) if isinstance(child, Tree) else child