class EventLoop:
"""event loop, handles IO"""
def __init__(self):
self.co_result = {}
self.selector = DefaultSelector()
def run_until_complete(self, coroutine):
self._continue(coroutine, None)
while coroutine not in self.co_result:
self._tick()
result = self.co_result[coroutine]
del self.co_result[coroutine]
return result
def run_forever(self):
while True:
self._tick()
def start_coroutine(self, co):
self._continue(co, None)
def _continue(self, co, awaited):
try:
cmd = co.send(awaited)
except StopIteration as e:
self.co_result[co] = e.value
else:
if isinstance(cmd, EventLoopCommand.StreamAccept):
self._register(EVENT_READ, cmd.sock, (cmd, co))
elif isinstance(cmd, EventLoopCommand.StreamConnect):
try:
cmd.sock.connect(cmd.addr)
except BlockingIOError:
self._register(EVENT_WRITE, cmd.sock, (cmd, co))
except:
raise
else:
# connect finished immediately
self._continue(co, None)
elif isinstance(cmd, EventLoopCommand.StreamRecv):
self._register(EVENT_READ, cmd.sock, (cmd, co))
elif isinstance(cmd, EventLoopCommand.StreamSend):
self._register(EVENT_WRITE, cmd.sock, (cmd, co))
else:
raise RuntimeError('unknown cmd')
def _register(self, event, sock, data):
try:
key = self.selector.get_key(sock)
except KeyError:
self.selector.register(sock, event, {event: data})
else:
assert key.events & event == 0
key.data[event] = data
self.selector.modify(sock, key.events | event, key.data)
def _unregister(self, event, sock):
key = self.selector.get_key(sock)
assert key.events & event != 0
new_mask = key.events & ~event
del key.data[event]
if new_mask == 0:
self.selector.unregister(sock)
else:
self.selector.modify(sock, new_mask, key.data)
def _tick(self):
for key, mask in self.selector.select(1.0):
if mask & EVENT_WRITE != 0:
cmd, co = key.data[EVENT_WRITE]
self._unregister(EVENT_WRITE, cmd.sock)
if isinstance(cmd, EventLoopCommand.StreamConnect):
self._continue(co, cmd.sock.connect(cmd.addr))
elif isinstance(cmd, EventLoopCommand.StreamSend):
self._continue(co, cmd.sock.send(cmd.buf))
else:
raise RuntimeError('unknown cmd')
elif mask & EVENT_READ != 0:
cmd, co = key.data[EVENT_READ]
self._unregister(EVENT_READ, cmd.sock)
if isinstance(cmd, EventLoopCommand.StreamAccept):
self._continue(co, cmd.sock.accept())
elif isinstance(cmd, EventLoopCommand.StreamRecv):
self._continue(co, cmd.sock.recv(cmd.buf_len))
else:
raise RuntimeError('unknown cmd')
