Extension to python-socketio server, allowing to await particular messages from clients

better_server.py

import asyncio

import socketio


class AsynchronousSocketServer(socketio.AsyncServer):

    def __init__(self, *args, **kwargs):
        kwargs["async_handlers"] = True
        super().__init__(*args, **kwargs)
        self._handlers = {}
        self._futures = {}

    def event_from_client(self, event, client, timeout=None):
        """
        Await this to continue execution on particular event from particular client.
        """
        idx = (event, client.id)
        if idx in self._futures:
            raise RuntimeError("already waiting for this message from this client")
        if event not in self._handlers:
            self._make_handler(event)
        future = asyncio.Future()
        self._futures[idx] = future
        if timeout is not None:
            self.start_background_task(self._set_timeout, idx, future, timeout)
        return future

    def _make_handler(self, event):
        self.logger.info(f"Registering handler for event {event!r}")

        def handler(sid, data=None):
            self.logger.debug(f"Got message {event!r} from client {sid}")
            idx = (event, sid)
            if idx in self._futures:
                future = self._futures.pop(idx)
                future.set_result(data)
            else:
                self.logger.warning(
                    "Got message {event!r} from client {sid}, but it will be ignored")
        self.on(event, handler)
        self._handlers[event] = handler

    async def _set_timeout(self, idx, future, timeout):
        await self.sleep(timeout)
        if not future.done():
            del self._futures[idx]
            future.set_exception(asyncio.TimeoutError(f"timeout of {timeout} seconds exceeded"))