extmod/uasyncio: Add new implementation of uasyncio module.

This commit adds a completely new implementation of the uasyncio module.
The aim of this version (compared to the original one in micropython-lib)
is to be more compatible with CPython's asyncio module, so that one can
more easily write code that runs under both MicroPython and CPython (and
reuse CPython asyncio libraries, follow CPython asyncio tutorials, etc).
Async code is not easy to write and any knowledge users already have from
CPython asyncio should transfer to uasyncio without effort, and vice versa.

The implementation here attempts to provide good compatibility with
CPython's asyncio while still being "micro" enough to run where MicroPython
runs. This follows the general philosophy of MicroPython itself, to make it
feel like Python.

The main change is to use a Task object for each coroutine.  This allows
more flexibility to queue tasks in various places, eg the main run loop,
tasks waiting on events, locks or other tasks.  It no longer requires
pre-allocating a fixed queue size for the main run loop.

A pairing heap is used to queue Tasks.

It's currently implemented in pure Python, separated into components with
lazy importing for optional components.  In the future parts of this
implementation can be moved to C to improve speed and reduce memory usage.
But the aim is to maintain a pure-Python version as a reference version.

1 files changed, 141 insertions, 0 deletions

diff --git a/extmod/uasyncio/stream.py b/extmod/uasyncio/stream.py
new file mode 100644
index 0000000000..7803ac4bfa
--- /dev/null
+++ b/extmod/uasyncio/stream.py
@@ -0,0 +1,141 @@
+# MicroPython uasyncio module
+# MIT license; Copyright (c) 2019-2020 Damien P. George
+
+from . import core
+
+
+class Stream:
+    def __init__(self, s, e={}):
+        self.s = s
+        self.e = e
+        self.out_buf = b""
+
+    def get_extra_info(self, v):
+        return self.e[v]
+
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        await self.close()
+
+    def close(self):
+        pass
+
+    async def wait_closed(self):
+        # TODO yield?
+        self.s.close()
+
+    async def read(self, n):
+        yield core._io_queue.queue_read(self.s)
+        return self.s.read(n)
+
+    async def readline(self):
+        l = b""
+        while True:
+            yield core._io_queue.queue_read(self.s)
+            l2 = self.s.readline()  # may do multiple reads but won't block
+            l += l2
+            if not l2 or l[-1] == 10:  # \n (check l in case l2 is str)
+                return l
+
+    def write(self, buf):
+        self.out_buf += buf
+
+    async def drain(self):
+        mv = memoryview(self.out_buf)
+        off = 0
+        while off < len(mv):
+            yield core._io_queue.queue_write(self.s)
+            ret = self.s.write(mv[off:])
+            if ret is not None:
+                off += ret
+        self.out_buf = b""
+
+
+# Create a TCP stream connection to a remote host
+async def open_connection(host, port):
+    from uerrno import EINPROGRESS
+    import usocket as socket
+
+    ai = socket.getaddrinfo(host, port)[0]  # TODO this is blocking!
+    s = socket.socket()
+    s.setblocking(False)
+    ss = Stream(s)
+    try:
+        s.connect(ai[-1])
+    except OSError as er:
+        if er.args[0] != EINPROGRESS:
+            raise er
+    yield core._io_queue.queue_write(s)
+    return ss, ss
+
+
+# Class representing a TCP stream server, can be closed and used in "async with"
+class Server:
+    async def __aenter__(self):
+        return self
+
+    async def __aexit__(self, exc_type, exc, tb):
+        self.close()
+        await self.wait_closed()
+
+    def close(self):
+        self.task.cancel()
+
+    async def wait_closed(self):
+        await self.task
+
+    async def _serve(self, cb, host, port, backlog):
+        import usocket as socket
+
+        ai = socket.getaddrinfo(host, port)[0]  # TODO this is blocking!
+        s = socket.socket()
+        s.setblocking(False)
+        s.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
+        s.bind(ai[-1])
+        s.listen(backlog)
+        self.task = core.cur_task
+        # Accept incoming connections
+        while True:
+            try:
+                yield core._io_queue.queue_read(s)
+            except core.CancelledError:
+                # Shutdown server
+                s.close()
+                return
+            try:
+                s2, addr = s.accept()
+            except:
+                # Ignore a failed accept
+                continue
+            s2.setblocking(False)
+            s2s = Stream(s2, {"peername": addr})
+            core.create_task(cb(s2s, s2s))
+
+
+# Helper function to start a TCP stream server, running as a new task
+# TODO could use an accept-callback on socket read activity instead of creating a task
+async def start_server(cb, host, port, backlog=5):
+    s = Server()
+    core.create_task(s._serve(cb, host, port, backlog))
+    return s
+
+
+################################################################################
+# Legacy uasyncio compatibility
+
+
+async def stream_awrite(self, buf, off=0, sz=-1):
+    if off != 0 or sz != -1:
+        buf = memoryview(buf)
+        if sz == -1:
+            sz = len(buf)
+        buf = buf[off : off + sz]
+    self.write(buf)
+    await self.drain()
+
+
+Stream.aclose = Stream.wait_closed
+Stream.awrite = stream_awrite
+Stream.awritestr = stream_awrite  # TODO explicitly convert to bytes?