extmod/uasyncio: Add asyncio.wait_for_ms function.

Fixes issue #6107.

2 files changed, 9 insertions, 4 deletions

diff --git a/extmod/uasyncio/__init__.py b/extmod/uasyncio/__init__.py
index da8b58061..08f924cf2 100644
--- a/extmod/uasyncio/__init__.py
+++ b/extmod/uasyncio/__init__.py
@@ -7,6 +7,7 @@ __version__ = (3, 0, 0)
 
 _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
index 7a4bddf25..6e1305c94 100644
--- 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 @@ async def wait_for(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)):