Repository: lark-parser/lark
Pull request #1000 (2021-09-15 13:48:52)

Fix for smart_decorator
- [x] new test to catch edge cases.
- [x] find all edge cases? (@erezsh and @ornariece I would like to know if you find more edge cases)
- [x] rework `smart_decorator`


Patch:
diff --git a/examples/advanced/dynamic_complete.py b/examples/advanced/dynamic_complete.py
index 48d5048..386936d 100644
--- a/examples/advanced/dynamic_complete.py
+++ b/examples/advanced/dynamic_complete.py
@@ -36,8 +36,8 @@ GRAMMAR = r"""
       | "false"            -> false
       | "null"             -> null
 
-array  : "[" [value ("," value)*] "]"
-object : "{" [pair ("," pair)*] "}"
+array  : "[" (value ("," value)*)? "]"
+object : "{" (pair ("," pair)*)? "}"
 pair   : string ":" value
 
 string: STRING
@@ -140,5 +140,4 @@ for example in EXAMPLES:
     tree = parser.parse(example)
     tree = RemoveAmbiguities().transform(tree)
     result = TreeToJson().transform(tree)
-    print('-' * 100)
     pprint(result)
diff --git a/lark/utils.py b/lark/utils.py
index db6cb7c..583caa2 100644
--- a/lark/utils.py
+++ b/lark/utils.py
@@ -108,29 +108,6 @@ class SerializeMemoizer(Serialize):
         return _deserialize(data, namespace, memo)
 
 
-
-import types
-from functools import wraps, partial
-
-
-def smart_decorator(f, create_decorator):
-    if isinstance(f, types.FunctionType):
-        return wraps(f)(create_decorator(f, True))
-
-    elif isinstance(f, (type, types.BuiltinFunctionType)):
-        return wraps(f)(create_decorator(f, False))
-
-    elif isinstance(f, types.MethodType):
-        return wraps(f)(create_decorator(f.__func__, True))
-
-    elif isinstance(f, partial):
-        # wraps does not work for partials in 2.7: https://bugs.python.org/issue3445
-        return wraps(f.func)(create_decorator(lambda *args, **kw: f(*args[1:], **kw), True))
-
-    else:
-        return create_decorator(f.__func__.__call__, True)
-
-
 try:
     import regex  # type: ignore
 except ImportError:
diff --git a/lark/visitors.py b/lark/visitors.py
index 028d58f..03dc740 100644
--- a/lark/visitors.py
+++ b/lark/visitors.py
@@ -1,8 +1,8 @@
-from typing import TypeVar, Tuple, List, Callable, Generic, Type, Union, Optional
+from typing import TypeVar, Tuple, List, Callable, Generic, Type, Union, Optional, Any
 from abc import ABC
-from functools import wraps
+from functools import wraps, update_wrapper
 
-from .utils import smart_decorator, combine_alternatives
+from .utils import combine_alternatives
 from .tree import Tree
 from .exceptions import VisitError, GrammarError
 from .lexer import Token
@@ -28,7 +28,7 @@ class _Decoratable:
     "Provides support for decorating methods with @v_args"
 
     @classmethod
-    def _apply_decorator(cls, decorator, **kwargs):
+    def _apply_v_args(cls, visit_wrapper):
         mro = getmro(cls)
         assert mro[0] is cls
         libmembers = {name for _cls in mro[1:] for name, _ in getmembers(_cls)}
@@ -41,11 +41,10 @@ class _Decoratable:
                 continue
 
             # Skip if v_args already applied (at the function level)
-            if hasattr(cls.__dict__[name], 'vargs_applied') or hasattr(value, 'vargs_applied'):
+            if isinstance(cls.__dict__[name], _VArgsWrapper):
                 continue
 
-            static = isinstance(cls.__dict__[name], (staticmethod, classmethod))
-            setattr(cls, name, decorator(value, static=static, **kwargs))
+            setattr(cls, name, _VArgsWrapper(cls.__dict__[name], visit_wrapper))
         return cls
 
     def __class_getitem__(cls, _):
@@ -402,46 +401,47 @@ def visit_children_decor(func: _InterMethod) -> _InterMethod:
 
 # Decorators
 
-def _apply_decorator(obj, decorator, **kwargs):
+def _apply_v_args(obj, visit_wrapper):
     try:
-        _apply = obj._apply_decorator
+        _apply = obj._apply_v_args
     except AttributeError:
-        return decorator(obj, **kwargs)
+        return _VArgsWrapper(obj, visit_wrapper)
     else:
-        return _apply(decorator, **kwargs)
+        return _apply(visit_wrapper)
 
 
-def _inline_args__func(func):
-    @wraps(func)
-    def create_decorator(_f, with_self):
-        if with_self:
-            def f(self, children):
-                return _f(self, *children)
-        else:
-            def f(self, children):
-                return _f(*children)
-        return f
-
-    return smart_decorator(func, create_decorator)
+class _VArgsWrapper:
+    """
+    A wrapper around a Callable. It delegates `__call__` to the Callable.
+    If the Callable has a `__get__`, that is also delegate and the resulting function is wrapped.
+    Otherwise, we use the original function mirroring the behaviour without a __get__.
+    We also have the visit_wrapper attribute to be used by Transformers.
+    """
+    def __init__(self, func: Callable, visit_wrapper: Callable[[Callable, str, list, Any], Any]):
+        if isinstance(func, _VArgsWrapper):
+            func = func.base_func
+        self.base_func = func
+        self.visit_wrapper = visit_wrapper
+        update_wrapper(self, func)
 
+    def __call__(self, *args, **kwargs):
+        return self.base_func(*args, **kwargs)
 
-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):
-                return _f(self, *args, **kwargs)
+    def __get__(self, instance, owner=None):
+        try:
+            g = self.base_func.__get__
+        except AttributeError:
+            return self
         else:
-            def f(self, *args, **kwargs):
-                return _f(*args, **kwargs)
-        return f
+            return _VArgsWrapper(g(instance, owner), self.visit_wrapper)
 
-    if static:
-        f = wraps(func)(create_decorator(func, False))
-    else:
-        f = smart_decorator(func, create_decorator)
-    f.vargs_applied = True
-    f.visit_wrapper = visit_wrapper
-    return f
+    def __set_name__(self, owner, name):
+        try:
+            f = self.base_func.__set_name__
+        except AttributeError:
+            return
+        else:
+            f(owner, name)
 
 
 def _vargs_inline(f, _data, children, _meta):
@@ -454,7 +454,7 @@ def _vargs_tree(f, data, children, meta):
     return f(Tree(data, children, meta))
 
 
-def v_args(inline: bool=False, meta: bool=False, tree: bool=False, wrapper: Optional[Callable]=None) -> Callable[[_DECORATED], _DECORATED]:
+def v_args(inline: bool = False, meta: bool = False, tree: bool = False, wrapper: Optional[Callable] = None) -> Callable[[_DECORATED], _DECORATED]:
     """A convenience decorator factory for modifying the behavior of user-supplied visitor methods.
 
     By default, callback methods of transformers/visitors accept one argument - a list of the node's children.
@@ -505,7 +505,7 @@ def v_args(inline: bool=False, meta: bool=False, tree: bool=False, wrapper: Opti
         func = wrapper
 
     def _visitor_args_dec(obj):
-        return _apply_decorator(obj, _visitor_args_func_dec, visit_wrapper=func)
+        return _apply_v_args(obj, func)
     return _visitor_args_dec
 
 


Test patch:
diff --git a/tests/test_nearley/nearley b/tests/test_nearley/nearley
index 3268316..a46b374 160000
--- a/tests/test_nearley/nearley
+++ b/tests/test_nearley/nearley
@@ -1,1 +1,1 @@
-Subproject commit 326831689826cb1b9a4d21d1ce0d5db9278e9636
+Subproject commit a46b37471db486db0f6e1ce6a2934fb238346b44
diff --git a/tests/test_tools.py b/tests/test_tools.py
index fd42b1c..80e69b0 100644
--- a/tests/test_tools.py
+++ b/tests/test_tools.py
@@ -1,18 +1,12 @@
 from __future__ import absolute_import, print_function
 
-import sys
 from unittest import TestCase, main
 
 from lark import Lark
 from lark.tree import Tree
 from lark.tools import standalone
 
-try:
-    from StringIO import StringIO
-except ImportError:
-    from io import StringIO
-
-
+from io import StringIO
 
 
 class TestStandalone(TestCase):
diff --git a/tests/test_trees.py b/tests/test_trees.py
index 82bf6c9..36be73f 100644
--- a/tests/test_trees.py
+++ b/tests/test_trees.py
@@ -1,6 +1,8 @@
 from __future__ import absolute_import
 
 import unittest
+from functools import partial, reduce, partialmethod
+from operator import add, mul
 from unittest import TestCase
 import copy
 import pickle
@@ -171,6 +173,145 @@ class TestTrees(TestCase):
         x = MyTransformer().transform( Tree('hello', [2]))
         self.assertEqual(x, 'hello')
 
+    def test_smart_decorator(self):
+        class OtherClass:
+            @staticmethod
+            def ab_staticmethod(a, b):
+                return (a, b)
+
+            @classmethod
+            def ab_classmethod(cls, a, b):
+                assert cls is OtherClass, cls
+                return (a, b)
+
+            def ab_method(self, a, b):
+                assert isinstance(self, OtherClass), self
+                return (a, b)
+
+        @v_args(meta=True)
+        class OtherTransformer(Transformer):
+            @staticmethod
+            def ab_staticmethod(meta, children):
+                return tuple(children)
+
+            @classmethod
+            def ab_classmethod(cls, meta, children):
+                assert cls is OtherTransformer, cls
+                return tuple(children)
+
+            def ab_method(self, meta, children):
+                assert isinstance(self, OtherTransformer), self
+                return tuple(children)
+
+        class CustomCallable:
+            def __call__(self, *args, **kwargs):
+                assert isinstance(self, CustomCallable)
+                return args
+
+        oc_instance = OtherClass()
+        ot_instance = OtherTransformer()
+
+        def ab_for_partialmethod(self, a, b):
+            assert isinstance(self, TestCls)
+            return a, b
+
+        @v_args(inline=True)
+        class TestCls(Transformer):
+            @staticmethod
+            def ab_staticmethod(a, b):
+                return (a, b)
+
+            @classmethod
+            def ab_classmethod(cls, a, b):
+                assert cls is TestCls
+                return (a, b)
+
+            def ab_method(self, a, b):
+                assert isinstance(self, TestCls)
+                return (a, b)
+
+            oc_class_ab_staticmethod = oc_instance.ab_staticmethod
+            oc_class_ab_classmethod = oc_instance.ab_classmethod
+
+            oc_ab_staticmethod = oc_instance.ab_staticmethod
+            oc_ab_classmethod = oc_instance.ab_classmethod
+            oc_ab_method = oc_instance.ab_method
+
+            ot_class_ab_staticmethod = ot_instance.ab_staticmethod
+            ot_class_ab_classmethod = ot_instance.ab_classmethod
+
+            ot_ab_staticmethod = ot_instance.ab_staticmethod
+            ot_ab_classmethod = ot_instance.ab_classmethod
+            ot_ab_method = ot_instance.ab_method
+
+            ab_partialmethod = partialmethod(ab_for_partialmethod, 1)
+            set_union = set(["a"]).union
+            static_add = staticmethod(add)
+            partial_reduce_mul = partial(reduce, mul)
+
+            custom_callable = CustomCallable()
+
+        test_instance = TestCls()
+        expected = {
+            "ab_classmethod": ([1, 2], (1, 2)),
+            "ab_staticmethod": ([1, 2], (1, 2)),
+            "ab_method": ([1, 2], (1, 2)),
+            "oc_ab_classmethod": ([1, 2], (1, 2)),
+            "oc_class_ab_classmethod": ([1, 2], (1, 2)),
+
+            # AFAIK, these two cases are impossible to deal with. `oc_instance.ab_staticmethod` returns an actual
+            # function object that is impossible to distinguish from a normally defined method.
+            # (i.e. `staticmethod(f).__get__(?, ?) is f` is True)
+            # "oc_ab_staticmethod": ([1, 2], (1, 2)),
+            # "oc_class_ab_staticmethod": ([1, 2], (1, 2)),
+
+            "oc_ab_method": ([1, 2], (1, 2)),
+            "ot_ab_classmethod": ([1, 2], (1, 2)),
+            "ot_class_ab_classmethod": ([1, 2], (1, 2)),
+
+            # Same as above
+            # "ot_ab_staticmethod": ([1, 2], (1, 2)),
+            # "ot_class_ab_staticmethod": ([1, 2], (1, 2)),
+
+            "ot_ab_method": ([1, 2], (1, 2)),
+            "ab_partialmethod": ([2], (1, 2)),
+            "custom_callable": ([1, 2], (1, 2)),
+            "set_union": ([["b"], ["c"]], {"a", "b", "c"}),
+            "static_add": ([1, 2], 3),
+            "partial_reduce_mul": ([[1, 2]], 2),
+        }
+        non_static = {"ab_method", "ab_partialmethod"}
+        for method_name, (children, expected_result) in expected.items():
+            not_inline = "ot" in method_name
+            result = test_instance.transform(Tree(method_name, children))
+            self.assertEqual(result, expected_result)
+
+            if not_inline:
+                result = getattr(test_instance, method_name)(None, children)
+            else:
+                result = getattr(test_instance, method_name)(*children)
+            self.assertEqual(result, expected_result)
+
+            if method_name not in non_static:
+                if not_inline:
+                    result = getattr(TestCls, method_name)(None, children)
+                else:
+                    result = getattr(TestCls, method_name)(*children)
+                self.assertEqual(result, expected_result)
+
+    def test_vargs_set_name(self):
+        # Test with cached_property if available. That actually uses __set_name__
+        prop = getattr(functools, "cached_property", property)
+
+        class T(Transformer):
+            @v_args(inline=True)
+            @prop  # Not sure why you would ever want to use a property here, but we support it
+            def test(self):
+                return lambda a, b: (self, a, b)
+
+        t = T()
+        self.assertEqual(t.transform(Tree("test", [1, 2])), (t, 1, 2))
+
     def test_inline_static(self):
         @v_args(inline=True)
         class T(Transformer):
@@ -213,7 +354,6 @@ class TestTrees(TestCase):
         res = T().transform(tree)
         assert res.children == ["@TEST1!", "test2!"]
 
-
     def test_discard(self):
         class MyTransformer(Transformer):
             def a(self, args):