# 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 readexactly(self, n):

r = b""

while n:

yield 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):

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""

# 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):

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?