py/compile: Handle return/break/continue correctly in async with.

Before this patch the context manager's __aexit__() method would not be
executed if a return/break/continue statement was used to exit an async
with block.  async with now has the same semantics as normal with.

The fix here applies purely to the compiler, and does not modify the
runtime at all. It might (eventually) be better to define new bytecode(s)
to handle async with (and maybe other async constructs) in a cleaner, more
efficient way.

One minor drawback with addressing this issue purely in the compiler is
that it wasn't possible to get 100% CPython semantics.  The thing that is
different here to CPython is that the __aexit__ method is not looked up in
the context manager until it is needed, which is after the body of the
async with statement has executed.  So if a context manager doesn't have
__aexit__ then CPython raises an exception before the async with is
executed, whereas uPy will raise it after it is executed.  Note that
__aenter__ is looked up at the beginning in uPy because it needs to be
called straightaway, so if the context manager isn't a context manager then
it'll still raise an exception at the same location as CPython.  The only
difference is if the context manager has the __aenter__ method but not the
__aexit__ method, then in that case uPy has different behaviour.  But this
is a very minor, and acceptable, difference.

5 files changed, 140 insertions, 1 deletions

diff --git a/tests/basics/async_with_break.py b/tests/basics/async_with_break.py
new file mode 100644
index 000000000..39bcbccb0
--- /dev/null
+++ b/tests/basics/async_with_break.py
@@ -0,0 +1,59 @@
+# test async with, escaped by a break
+
+class AContext:
+    async def __aenter__(self):
+        print('enter')
+        return 1
+    async def __aexit__(self, exc_type, exc, tb):
+        print('exit', exc_type, exc)
+
+async def f1():
+    while 1:
+        async with AContext():
+            print('body')
+            break
+            print('no 1')
+        print('no 2')
+
+o = f1()
+try:
+    print(o.send(None))
+except StopIteration:
+    print('finished')
+
+async def f2():
+    while 1:
+        try:
+            async with AContext():
+                print('body')
+                break
+                print('no 1')
+        finally:
+            print('finally')
+        print('no 2')
+
+o = f2()
+try:
+    print(o.send(None))
+except StopIteration:
+    print('finished')
+
+async def f3():
+    while 1:
+        try:
+            try:
+                async with AContext():
+                    print('body')
+                    break
+                    print('no 1')
+            finally:
+                print('finally inner')
+        finally:
+            print('finally outer')
+        print('no 2')
+
+o = f3()
+try:
+    print(o.send(None))
+except StopIteration:
+    print('finished')
diff --git a/tests/basics/async_with_break.py.exp b/tests/basics/async_with_break.py.exp
new file mode 100644
index 000000000..d077a88fa
--- /dev/null
+++ b/tests/basics/async_with_break.py.exp
@@ -0,0 +1,15 @@
+enter
+body
+exit None None
+finished
+enter
+body
+exit None None
+finally
+finished
+enter
+body
+exit None None
+finally inner
+finally outer
+finished
diff --git a/tests/basics/async_with_return.py b/tests/basics/async_with_return.py
new file mode 100644
index 000000000..9af88b839
--- /dev/null
+++ b/tests/basics/async_with_return.py
@@ -0,0 +1,50 @@
+# test async with, escaped by a return
+
+class AContext:
+    async def __aenter__(self):
+        print('enter')
+        return 1
+    async def __aexit__(self, exc_type, exc, tb):
+        print('exit', exc_type, exc)
+
+async def f1():
+    async with AContext():
+        print('body')
+        return
+
+o = f1()
+try:
+    o.send(None)
+except StopIteration:
+    print('finished')
+
+async def f2():
+    try:
+        async with AContext():
+            print('body')
+            return
+    finally:
+        print('finally')
+
+o = f2()
+try:
+    o.send(None)
+except StopIteration:
+    print('finished')
+
+async def f3():
+    try:
+        try:
+            async with AContext():
+                print('body')
+                return
+        finally:
+            print('finally inner')
+    finally:
+        print('finally outer')
+
+o = f3()
+try:
+    o.send(None)
+except StopIteration:
+    print('finished')
diff --git a/tests/basics/async_with_return.py.exp b/tests/basics/async_with_return.py.exp
new file mode 100644
index 000000000..d077a88fa
--- /dev/null
+++ b/tests/basics/async_with_return.py.exp
@@ -0,0 +1,15 @@
+enter
+body
+exit None None
+finished
+enter
+body
+exit None None
+finally
+finished
+enter
+body
+exit None None
+finally inner
+finally outer
+finished
diff --git a/tests/run-tests b/tests/run-tests
index cfd7c4037..e1b594edf 100755
--- a/tests/run-tests
+++ b/tests/run-tests
@@ -338,7 +338,7 @@ def run_tests(pyb, tests, args, base_path="."):
     if args.emit == 'native':
         skip_tests.update({'basics/%s.py' % t for t in 'gen_yield_from gen_yield_from_close gen_yield_from_ducktype gen_yield_from_exc gen_yield_from_executing gen_yield_from_iter gen_yield_from_send gen_yield_from_stopped gen_yield_from_throw gen_yield_from_throw2 gen_yield_from_throw3 generator1 generator2 generator_args generator_close generator_closure generator_exc generator_pend_throw generator_return generator_send'.split()}) # require yield
         skip_tests.update({'basics/%s.py' % t for t in 'bytes_gen class_store_class globals_del string_join'.split()}) # require yield
-        skip_tests.update({'basics/async_%s.py' % t for t in 'def await await2 for for2 with with2'.split()}) # require yield
+        skip_tests.update({'basics/async_%s.py' % t for t in 'def await await2 for for2 with with2 with_break with_return'.split()}) # require yield
         skip_tests.update({'basics/%s.py' % t for t in 'try_reraise try_reraise2'.split()}) # require raise_varargs
         skip_tests.update({'basics/%s.py' % t for t in 'with_break with_continue with_return'.split()}) # require complete with support
         skip_tests.add('basics/array_construct2.py') # requires generators