# SPDX-FileCopyrightText: 2019-2020 Damien P. George # # SPDX-License-Identifier: MIT # # MicroPython uasyncio module # MIT license; Copyright (c) 2019-2020 Damien P. George # # This code comes from MicroPython, and has not been run through black or pylint there. # Altering these files significantly would make merging difficult, so we will not use # pylint or black. # pylint: skip-file # fmt: off """ Streams ======= """ from . import core class Stream: """This represents a TCP stream connection. To minimise code this class implements both a reader and a writer, and both ``StreamReader`` and ``StreamWriter`` alias to this class. """ def __init__(self, s, e={}): self.s = s self.e = e self.out_buf = b"" def get_extra_info(self, v): """Get extra information about the stream, given by *v*. The valid values for *v* are: ``peername``. """ 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): """Wait for the stream to close. This is a coroutine. """ # TODO yield? self.s.close() async def read(self, n): """Read up to *n* bytes and return them. This is a coroutine. """ await core._io_queue.queue_read(self.s) return self.s.read(n) async def readinto(self, buf): """Read up to n bytes into *buf* with n being equal to the length of *buf* Return the number of bytes read into *buf* This is a coroutine, and a MicroPython extension. """ await core._io_queue.queue_read(self.s) return self.s.readinto(buf) async def readexactly(self, n): """Read exactly *n* bytes and return them as a bytes object. Raises an ``EOFError`` exception if the stream ends before reading *n* bytes. This is a coroutine. """ r = b"" while n: await core._io_queue.queue_read(self.s) r2 = self.s.read(n) if r2 is not None: if not len(r2): raise EOFError r += r2 n -= len(r2) return r async def readline(self): """Read a line and return it. This is a coroutine. """ l = b"" while True: await 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): """Accumulated *buf* to the output buffer. The data is only flushed when `Stream.drain` is called. It is recommended to call `Stream.drain` immediately after calling this function. """ if not self.out_buf: # Try to write immediately to the underlying stream. ret = self.s.write(buf) if ret == len(buf): return if ret is not None: buf = buf[ret:] self.out_buf += buf async def drain(self): """Drain (write) all buffered output data out to the stream. This is a coroutine. """ mv = memoryview(self.out_buf) off = 0 while off < len(mv): await core._io_queue.queue_write(self.s) ret = self.s.write(mv[off:]) if ret is not None: off += ret self.out_buf = b"" # Stream can be used for both reading and writing to save code size StreamReader = Stream StreamWriter = Stream # Create a TCP stream connection to a remote host async def open_connection(host, port): """Open a TCP connection to the given *host* and *port*. The *host* address will be resolved using `socket.getaddrinfo`, which is currently a blocking call. Returns a pair of streams: a reader and a writer stream. Will raise a socket-specific ``OSError`` if the host could not be resolved or if the connection could not be made. This is a coroutine. """ from uerrno import EINPROGRESS import usocket as socket ai = socket.getaddrinfo(host, port, 0, socket.SOCK_STREAM)[0] # TODO this is blocking! s = socket.socket(ai[0], ai[1], ai[2]) s.setblocking(False) ss = Stream(s) try: s.connect(ai[-1]) except OSError as er: if er.errno != EINPROGRESS: raise er await 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: """This represents the server class returned from `start_server`. It can be used in an ``async with`` statement to close the server upon exit. """ async def __aenter__(self): return self async def __aexit__(self, exc_type, exc, tb): self.close() await self.wait_closed() def close(self): """Close the server.""" self.task.cancel() async def wait_closed(self): """Wait for the server to close. This is a coroutine. """ await self.task async def _serve(self, s, cb): # Accept incoming connections while True: try: await 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): """Start a TCP server on the given *host* and *port*. The *cb* callback will be called with incoming, accepted connections, and be passed 2 arguments: reader writer streams for the connection. Returns a `Server` object. This is a coroutine. """ import usocket as socket # Create and bind server socket. host = 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(host[-1]) s.listen(backlog) # Create and return server object and task. srv = Server() srv.task = core.create_task(srv._serve(s, cb)) return srv ################################################################################ # 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?