enh: curry macro: detect manual curry, skip injecting a curry when detected

Technologicat · Technologicat · commit 4625951d8149 · 2019-02-11T11:38:58.000+02:00
diff --git a/unpythonic/syntax/curry.py b/unpythonic/syntax/curry.py
@@ -1,41 +1,44 @@
 # -*- coding: utf-8 -*-
 """Automatic currying. Transforms both function definitions and calls."""
 
-from ast import Call, Lambda, FunctionDef, With, withitem
+from ast import Call, Lambda, FunctionDef
 from .astcompat import AsyncFunctionDef
 
 from macropy.core.quotes import macros, ast_literal
 from macropy.core.hquotes import macros, hq
 from macropy.core.walkers import Walker
 
-from .util import suggest_decorator_index
-
-from ..dynassign import dyn
+from .util import suggest_decorator_index, isx, make_isxpred, has_curry, \
+                  sort_lambda_decorators
 
 # CAUTION: unpythonic.syntax.lambdatools.namedlambda depends on the exact names
 # "curryf" and "currycall" to detect an auto-curried expression with a final lambda.
 from ..fun import curry as curryf, _currycall as currycall
 
+_iscurry = make_isxpred("curry")
+
 def curry(block_body):
     @Walker
-    def transform_call(tree, *, stop, **kw):  # technically a node containing the current subtree
+    def transform(tree, *, hascurry, set_ctx, stop, **kw):
         if type(tree) is Call:
-            tree.args = [tree.func] + tree.args
-            tree.func = hq[currycall]
+            if has_curry(tree):  # detect decorated lambda with manual curry
+                set_ctx(hascurry=True)  # the lambda inside the curry(...) is the next Lambda node we will descend into.
+            if not isx(tree.func, _iscurry):
+                tree.args = [tree.func] + tree.args
+                tree.func = hq[currycall]
         elif type(tree) in (FunctionDef, AsyncFunctionDef):
-            # TODO: detect there's no curry already.
-            k = suggest_decorator_index("curry", tree.decorator_list)
-            if k is not None:
-                tree.decorator_list.insert(k, hq[curryf])
-            else:  # couldn't determine insert position; just plonk it at the end and hope for the best
-                tree.decorator_list.append(hq[curryf])
+            if not any(isx(item, _iscurry) for item in tree.decorator_list):  # no manual curry already
+                k = suggest_decorator_index("curry", tree.decorator_list)
+                if k is not None:
+                    tree.decorator_list.insert(k, hq[curryf])
+                else:  # couldn't determine insert position; just plonk it at the end and hope for the best
+                    tree.decorator_list.append(hq[curryf])
         elif type(tree) is Lambda:
-            # This inserts curry() as the innermost "decorator", and the curry
-            # macro is meant to run last (after e.g. tco), so we're fine.
-            # TODO: detect there's no curry already.
-            tree = hq[curryf(ast_literal[tree])]
-            # don't recurse on the lambda we just moved, but recurse inside it.
-            stop()
-            tree.args[0].body = transform_call.recurse(tree.args[0].body)
+            if not hascurry:
+                tree = hq[curryf(ast_literal[tree])]
+                # don't recurse on the lambda we just moved, but recurse inside it.
+                stop()
+                tree.args[0].body = transform.recurse(tree.args[0].body, hascurry=False)
         return tree
-    return transform_call.recurse(block_body)
+    newbody = transform.recurse(block_body, hascurry=False)
+    return sort_lambda_decorators(newbody)
diff --git a/unpythonic/syntax/tailtools.py b/unpythonic/syntax/tailtools.py
@@ -26,7 +26,6 @@
                   suggest_decorator_index
 from .letdoutil import isdo, islet, ExpandedLetView, ExpandedDoView
 from .ifexprs import aif
-from .letdo import let
 
 from ..dynassign import dyn
 from ..it import uniqify
diff --git a/unpythonic/syntax/test/test_curry.py b/unpythonic/syntax/test/test_curry.py
@@ -7,6 +7,8 @@
 from ...fun import composerc as compose
 from ...llist import cons, nil, ll
 
+from macropy.tracing import macros, show_expanded
+
 def test():
     with curry:
         mymap = lambda f: foldr(compose(cons, f), nil)
@@ -49,4 +51,23 @@ def stuffinto(lst, x):
     stuffinto(lst, 5)
     assert lst == [1, 2, 3, 4, 5]
 
+    # should not insert an extra @curry even if we curry manually
+    # (convenience, for with-currying existing code)
+    with show_expanded:
+     with curry:
+        from unpythonic.fun import curry
+        @curry
+        def add3(a, b, c):
+            return a + b + c
+        assert add3(1)(2)(3) == 6
+
+        f = curry(lambda a, b, c: a + b + c)
+        assert f(1)(2)(3) == 6
+
+        from unpythonic.tco import trampolined, jump
+        from unpythonic.fun import withself
+        fact = trampolined(withself(curry(lambda self, n, acc=1:
+                                    acc if n == 0 else jump(self, n - 1, n * acc))))
+        assert fact(5) == 120
+
     print("All tests PASSED")
diff --git a/unpythonic/syntax/util.py b/unpythonic/syntax/util.py
@@ -229,12 +229,26 @@ def has_tco(tree, userlambdas=[]):
     Return value is ``True`` or ``False`` (depending on test result) if the
     test was applicable, and ``None`` if it was not applicable (no match on tree).
     """
+    return _has_deco(tco_decorators, tree, userlambdas)
+
+def has_curry(tree, userlambdas=[]):
+    """Return whether a FunctionDef or a decorated lambda has curry applied.
+
+    userlambdas: list of ``id(some_tree)``; when detecting a lambda,
+    only consider it if its id matches one of those in the list.
+
+    Return value is ``True`` or ``False`` (depending on test result) if the
+    test was applicable, and ``None`` if it was not applicable (no match on tree).
+    """
+    return _has_deco(["curry"], tree, userlambdas)
+
+def _has_deco(deconames, tree, userlambdas=[]):
     if type(tree) in (FunctionDef, AsyncFunctionDef):
-        return any(is_decorator(x, fname) for fname in tco_decorators for x in tree.decorator_list)
+        return any(is_decorator(x, fname) for fname in deconames for x in tree.decorator_list)
     elif is_decorated_lambda(tree, mode="any"):
         decorator_list, thelambda = destructure_decorated_lambda(tree)
         if (not userlambdas) or (id(thelambda) in userlambdas):
-            return any(is_lambda_decorator(x, fname) for fname in tco_decorators for x in decorator_list)
+            return any(is_lambda_decorator(x, fname) for fname in deconames for x in decorator_list)
     return None  # not applicable
 
 def sort_lambda_decorators(tree):
