Fast Python Poll/Epoll Server

PythonPollServer.py

import socket
import selectors

in_stream_selector = selectors.DefaultSelector()
out_stream_selector = selectors.DefaultSelector()

write_handlers = {}


class WriteHandler:
    def __init__(self, data: bytes, start=0):
        self.data = bytearray(data)
        self.start = start
        self.remain = len(data)

    def append(self, new_data):
        self.data.extend(new_data)

    def write(self, conn: socket.socket):
        send_len = conn.send(bytes(self.data[self.start:]))
        self.start += send_len
        self.remain -= send_len
        if self.remain == 0:
            return True
        else:
            return False

    def __call__(self, conn: socket.socket, mask: int):
        self.write(conn)


def accept(sock: socket.socket, mask: int):
    conn, addr = sock.accept()  # Should be ready
    print('accepted', conn, 'from', addr)
    conn.setblocking(False)
    in_stream_selector.register(conn, selectors.EVENT_READ, handle_request)


def handle_request(conn: socket.socket, mask: int):
    data = conn.recv(1024)  # Should be ready
    if data:
        if b'\x1b' in data:
            print('close conn by recv ESC')
            close_conn(conn)
            return
        print('echoing', repr(data), 'to', conn)
        if conn not in write_handlers:
            write_handlers[conn] = WriteHandler(data)
            out_stream_selector.register(conn, selectors.EVENT_WRITE, write_handlers[conn])
        else:
            handler: WriteHandler = write_handlers[conn]
            handler.append(data)
    else:
        print('closing by no data', conn)
        close_conn(conn)


def close_conn(conn):
    in_stream_selector.unregister(conn)
    if conn in write_handlers:
        write_handlers.pop(conn)
        out_stream_selector.unregister(conn)
    conn.close()


def main():
    server_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
    server_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
    server_socket.bind(('0.0.0.0', 8999))
    server_socket.listen()
    server_socket.setblocking(False)

    in_stream_selector.register(server_socket, selectors.EVENT_READ, data=accept)
    while True:
        ready_list = in_stream_selector.select()
        for key, mask in ready_list:
            # key type: selectors.SelectorKey
            callback = key.data
            callback(key.fileobj, mask)
        if len(write_handlers) > 0:
            ready_list = out_stream_selector.select()
            for key, mask in ready_list:
                conn = key.fileobj
                handler: WriteHandler = key.data
                handler(conn, mask)


if __name__ == '__main__':
    main()