From 783d53ecb77505d855f7c22cd918c82995224eb1 Mon Sep 17 00:00:00 2001
From: Simon Gomizelj <simon@vodik.xyz>
Date: Sat, 30 Dec 2017 17:25:26 -0500
Subject: [PATCH] Introduce with/a* and with/a expressions

---
 hy/compiler.py                        | 19 ++++++-----
 hy/core/macros.hy                     | 35 +++++++++++++------
 tests/native_tests/py35_only_tests.hy | 48 +++++++++++++++++++++++++++
 3 files changed, 84 insertions(+), 18 deletions(-)

diff --git a/hy/compiler.py b/hy/compiler.py
index e7119b4..8ab1adb 100755
--- a/hy/compiler.py
+++ b/hy/compiler.py
@@ -1309,17 +1309,19 @@ class HyASTCompiler(object):
         return ret + fn
 
     @builds("with*")
+    @builds("with/a*", iff=PY35)
     @checkargs(min=2)
     def compile_with_expression(self, expr):
-        expr.pop(0)  # with*
+        root = expr.pop(0)
 
         args = expr.pop(0)
         if not isinstance(args, HyList):
             raise HyTypeError(expr,
-                              "with expects a list, received `{0}'".format(
-                                  type(args).__name__))
+                              "{0} expects a list, received `{1}'".format(
+                                  root, type(args).__name__))
         if len(args) not in (1, 2):
-            raise HyTypeError(expr, "with needs [arg (expr)] or [(expr)]")
+            raise HyTypeError(expr,
+                              "{0} needs [arg (expr)] or [(expr)]".format(root))
 
         thing = None
         if len(args) == 2:
@@ -1338,10 +1340,11 @@ class HyASTCompiler(object):
             expr, targets=[name], value=asty.Name(
                 expr, id=ast_str("None"), ctx=ast.Load()))
 
-        the_with = asty.With(expr,
-                             context_expr=ctx.force_expr,
-                             optional_vars=thing,
-                             body=body.stmts)
+        node = asty.With if root == "with*" else asty.AsyncWith
+        the_with = node(expr,
+                        context_expr=ctx.force_expr,
+                        optional_vars=thing,
+                        body=body.stmts)
 
         if PY3:
             the_with.items = [ast.withitem(context_expr=ctx.force_expr,
diff --git a/hy/core/macros.hy b/hy/core/macros.hy
index 208fa22..0052a16 100644
--- a/hy/core/macros.hy
+++ b/hy/core/macros.hy
@@ -43,6 +43,19 @@ be associated in pairs."
                                                other-kvs))]))))
 
 
+(defmacro _with [node args &rest body]
+  (if (not (empty? args))
+    (do
+     (if (>= (len args) 2)
+       (do
+        (setv p1 (.pop args 0)
+              p2 (.pop args 0)
+              primary [p1 p2])
+        `(~node [~@primary] (_with ~node ~args ~@body)))
+       `(~node [~@args] ~@body)))
+    `(do ~@body)))
+
+
 (defmacro with [args &rest body]
   "Wrap execution of `body` within a context manager given as bracket `args`.
 
@@ -51,16 +64,18 @@ Shorthand for nested with* loops:
   (with* [x foo]
     (with* [y bar]
       baz))."
-  (if (not (empty? args))
-    (do
-     (if (>= (len args) 2)
-       (do
-        (setv p1 (.pop args 0)
-              p2 (.pop args 0)
-              primary [p1 p2])
-        `(with* [~@primary] (with ~args ~@body)))
-       `(with* [~@args] ~@body)))
-    `(do ~@body)))
+  `(_with with* ~args ~@body))
+
+
+(defmacro with/a [args &rest body]
+  "Wrap execution of `body` with/ain a context manager given as bracket `args`.
+
+Shorthand for nested with/a* loops:
+  (with/a [x foo y bar] baz) ->
+  (with/a* [x foo]
+    (with/a* [y bar]
+      baz))."
+  `(_with with/a* ~args ~@body))
 
 
 (defmacro cond [&rest branches]
diff --git a/tests/native_tests/py35_only_tests.hy b/tests/native_tests/py35_only_tests.hy
index 7e144d7..9081dac 100644
--- a/tests/native_tests/py35_only_tests.hy
+++ b/tests/native_tests/py35_only_tests.hy
@@ -43,3 +43,51 @@
     (await (sleep 0))
     [1 2 3])
   (assert (= (run-coroutine coro-test) [1 2 3])))
+
+
+(defclass AsyncWithTest []
+  (defn --init-- [self val]
+    (setv self.val val)
+    None)
+
+  (defn/a --aenter-- [self]
+    self.val)
+
+  (defn/a --aexit-- [self tyle value traceback]
+    (setv self.val None)))
+
+
+(defn test-single-with/a []
+  (run-coroutine
+    (fn/a []
+      (with/a [t (AsyncWithTest 1)]
+        (assert (= t 1))))))
+
+(defn test-two-with/a []
+  (run-coroutine
+    (fn/a []
+      (with/a [t1 (AsyncWithTest 1)
+               t2 (AsyncWithTest 2)]
+        (assert (= t1 1))
+        (assert (= t2 2))))))
+
+(defn test-thrice-with/a []
+  (run-coroutine
+    (fn/a []
+      (with/a [t1 (AsyncWithTest 1)
+               t2 (AsyncWithTest 2)
+               t3 (AsyncWithTest 3)]
+        (assert (= t1 1))
+        (assert (= t2 2))
+        (assert (= t3 3))))))
+
+(defn test-quince-with/a []
+  (run-coroutine
+    (fn/a []
+      (with/a [t1 (AsyncWithTest 1)
+               t2 (AsyncWithTest 2)
+               t3 (AsyncWithTest 3)
+               _ (AsyncWithTest 4)]
+        (assert (= t1 1))
+        (assert (= t2 2))
+        (assert (= t3 3))))))