1# test waiting within async with enter/exit functions
2
3try:
4    import usys as sys
5except ImportError:
6    import sys
7if sys.implementation.name == 'micropython':
8    # uPy allows normal generators to be awaitables
9    coroutine = lambda f: f
10else:
11    import types
12    coroutine = types.coroutine
13
14@coroutine
15def f(x):
16    print('f start:', x)
17    yield x + 1
18    yield x + 2
19    return x + 3
20
21class AContext:
22    async def __aenter__(self):
23        print('enter')
24        print('f returned:', await f(10))
25    async def __aexit__(self, exc_type, exc, tb):
26        print('exit', exc_type, exc)
27        print('f returned:', await f(20))
28
29async def coro():
30    async with AContext():
31        print('body start')
32        print('body f returned:', await f(30))
33        print('body end')
34
35o = coro()
36try:
37    while True:
38        print('coro yielded:', o.send(None))
39except StopIteration:
40    print('finished')
41