extmod/uasyncio: Add asyncio.wait_for_ms function.

dpgeorge · dpgeorge · commit a4c96fb3b0c5 · 2020-06-10T22:29:44.000+10:00
Fixes issue adafruit#6107.
diff --git a/docs/library/uasyncio.rst b/docs/library/uasyncio.rst
@@ -74,6 +74,12 @@ Additional functions
 
     This is a coroutine.
 
+.. function:: wait_for_ms(awaitable, timeout)
+
+    Similar to `wait_for` but *timeout* is an integer in milliseconds.
+
+    This is a coroutine, and a MicroPython extension.
+
 .. function:: gather(\*awaitables, return_exceptions=False)
 
     Run all *awaitables* concurrently.  Any *awaitables* that are not tasks are
diff --git a/extmod/uasyncio/__init__.py b/extmod/uasyncio/__init__.py
@@ -7,6 +7,7 @@
 
 _attrs = {
     "wait_for": "funcs",
+    "wait_for_ms": "funcs",
     "gather": "funcs",
     "Event": "event",
     "Lock": "lock",
diff --git a/extmod/uasyncio/funcs.py b/extmod/uasyncio/funcs.py
@@ -4,16 +4,16 @@
 from . import core
 
 
-async def wait_for(aw, timeout):
+async def wait_for(aw, timeout, sleep=core.sleep):
     aw = core._promote_to_task(aw)
     if timeout is None:
         return await aw
 
-    def cancel(aw, timeout):
-        await core.sleep(timeout)
+    def cancel(aw, timeout, sleep):
+        await sleep(timeout)
         aw.cancel()
 
-    cancel_task = core.create_task(cancel(aw, timeout))
+    cancel_task = core.create_task(cancel(aw, timeout, sleep))
     try:
         ret = await aw
     except core.CancelledError:
@@ -29,6 +29,10 @@ def cancel(aw, timeout):
     return ret
 
 
+def wait_for_ms(aw, timeout):
+    return wait_for(aw, timeout, core.sleep_ms)
+
+
 async def gather(*aws, return_exceptions=False):
     ts = [core._promote_to_task(aw) for aw in aws]
     for i in range(len(ts)):
diff --git a/tests/extmod/uasyncio_micropython.py b/tests/extmod/uasyncio_micropython.py
@@ -0,0 +1,37 @@
+# Test MicroPython extensions on CPython asyncio:
+# - sleep_ms
+# - wait_for_ms
+
+try:
+    import utime, uasyncio
+except ImportError:
+    print("SKIP")
+    raise SystemExit
+
+
+async def task(id, t):
+    print("task start", id)
+    await uasyncio.sleep_ms(t)
+    print("task end", id)
+    return id * 2
+
+
+async def main():
+    # Simple sleep_ms
+    t0 = utime.ticks_ms()
+    await uasyncio.sleep_ms(1)
+    print(utime.ticks_diff(utime.ticks_ms(), t0) < 100)
+
+    # When task finished before the timeout
+    print(await uasyncio.wait_for_ms(task(1, 5), 50))
+
+    # When timeout passes and task is cancelled
+    try:
+        print(await uasyncio.wait_for_ms(task(2, 50), 5))
+    except uasyncio.TimeoutError:
+        print("timeout")
+
+    print("finish")
+
+
+uasyncio.run(main())
diff --git a/tests/extmod/uasyncio_micropython.py.exp b/tests/extmod/uasyncio_micropython.py.exp
@@ -0,0 +1,7 @@
+True
+task start 1
+task end 1
+2
+task start 2
+timeout
+finish
