enh: prefix macro now works together with let and do

Technologicat · Technologicat · commit c45bc01d7a3a · 2018-11-01T14:40:09.000+02:00
diff --git a/macro_extras/main.py b/macro_extras/main.py
@@ -1031,6 +1031,30 @@ def g(*args, **kwargs):
         a, b = (q, b, a)  # pythonic swap in prefix syntax; must quote RHS
         assert a == 200 and b == 100
 
+        # prefix has no effect on the let binding syntax ((name0, value0), ...)
+        a = let((x, 42))[x << x + 1]
+        assert a == 43
+
+        # but the RHSs of the bindings are transformed normally
+        def double(x):
+            return 2*x
+        a = let((x, (double, 21)))[x << x + 1]
+        assert a == 43
+
+        # similarly, prefix leaves the "body tuple" of a do alone
+        # (syntax, not semantically a tuple), but recurses into it:
+        a = do[1, 2, 3]
+        assert a == 3
+        a = do[1, 2, (double, 3)]
+        assert a == 6
+
+        # the extra bracket syntax has no danger of confusion, as it's a list, not tuple
+        a = let((x, 3))[[
+                  1,
+                  2,
+                  (double, x)]]
+        assert a == 6
+
     # Introducing LisThEll:
     with prefix, curry:  # important: apply prefix first, then curry
         mymap = lambda f: (foldr, (compose, cons, f), nil)
diff --git a/unpythonic/syntax/prefix.py b/unpythonic/syntax/prefix.py
@@ -4,7 +4,7 @@
 Experimental, not for use in production code.
 """
 
-from ast import Name, Call, Tuple, Load
+from ast import Name, Call, Tuple, Load, Subscript, Index
 
 from macropy.core.quotes import macros, q, u, ast_literal
 from macropy.core.walkers import Walker
@@ -16,7 +16,34 @@ def prefix(block_body):
     isunquote = lambda tree: type(tree) is Name and tree.id == "u"
     iskwargs = lambda tree: type(tree) is Call and type(tree.func) is Name and tree.func.id == "kw"
     @Walker
-    def transform(tree, *, quotelevel, set_ctx, **kw):
+    def transform(tree, *, quotelevel, set_ctx, stop, **kw):
+        # Not tuples but syntax: leave alone the:
+        #  - bindings blocks of let, letseq, letrec, and the d*, b* variants
+        #  - subscript part of an explicit do[]
+        # but recurse inside them.
+        #
+        # let and do have not expanded yet when prefix runs (better that way!),
+        # so we can't use the (expanded-form) detectors islet, isdo.
+        if type(tree) is Call and type(tree.func) is Name and \
+           any(tree.func.id == x for x in ("let", "letseq", "letrec",
+                                           "dlet", "dletseq", "dletrec",
+                                           "blet", "bletseq", "bletrec")):
+            # let((x, 42))[...] appears as Subscript(value=Call(...), ...)
+            stop()
+            for binding in tree.args:  # TODO: kwargs support for let(x=42)[...] if implemented later
+                _, value = binding.elts  # leave name alone, recurse into value
+                binding.elts[1] = transform.recurse(value, quotelevel=quotelevel)
+            return tree
+        elif type(tree) is Subscript and type(tree.value) is Name and \
+           any(tree.value.id == x for x in ("do", "do0")) and \
+           type(tree.slice) is Index and type(tree.slice.value) is Tuple:
+            stop()
+            newelts = []
+            for expr in tree.slice.value.elts:
+                newelts.append(transform.recurse(expr, quotelevel=quotelevel))
+            tree.slice.value.elts = newelts
+            return tree
+        # general case
         if not (type(tree) is Tuple and type(tree.ctx) is Load):
             return tree
         op, *data = tree.elts
