aiohttp_rpc.server.websocket

Complexity

Total Complexity

14

Size/Duplication

Total Lines	101
Duplicated Lines	0 %

Importance

Changes

0

5 Methods

Rating	Name	Duplication	Size	Complexity
A	WsJsonRpcServer.on_shutdown()	0	7	2
A	WsJsonRpcServer.__init__()	0	9	1
A	WsJsonRpcServer._handle_ws_message()	0	21	5
A	WsJsonRpcServer._handle_ws_request()	0	33	3
A	WsJsonRpcServer.handle_http_request()	0	5	3

import asyncio
import json
import typing
import weakref

from aiohttp import http_websocket, web, web_ws

from .base import BaseJsonRpcServer
from .. import errors, protocol, utils


__all__ = (
    'WsJsonRpcServer',
)


class WsJsonRpcServer(BaseJsonRpcServer):
    rcp_websockets: weakref.WeakSet
    _json_response_handler: typing.Optional[typing.Callable] = None
    _background_tasks: typing.Set

    def __init__(self,
                 *args,
                 json_response_handler: typing.Optional[typing.Callable] = None,
                 **kwargs) -> None:
        super().__init__(*args, **kwargs)

        self.rcp_websockets = weakref.WeakSet()
        self._json_response_handler = json_response_handler
        self._background_tasks = set()

    async def handle_http_request(self, http_request: web.Request) -> web.StreamResponse:
        if http_request.method != 'GET' or http_request.headers.get('upgrade', '').lower() != 'websocket':
            raise web.HTTPMethodNotAllowed(method=http_request.method, allowed_methods=('GET',))

        return await self._handle_ws_request(http_request)

    async def on_shutdown(self, app: web.Application) -> None:
        # https://docs.aiohttp.org/en/stable/web_advanced.html#graceful-shutdown

        for ws in self.rcp_websockets:
            await ws.close(code=http_websocket.WSCloseCode.GOING_AWAY, message='Server shutdown')

        self.rcp_websockets.clear()

    async def _handle_ws_request(self, http_request: web.Request) -> web_ws.WebSocketResponse:
        from aiohttp_rpc import WsJsonRpcClient

        ws_connect = web_ws.WebSocketResponse()
        await ws_connect.prepare(http_request)

        self.rcp_websockets.add(ws_connect)

        ws_rpc_client = WsJsonRpcClient(ws_connect=ws_connect)

        ws_msg: http_websocket.WSMessage

        async for ws_msg in ws_connect:
            if ws_msg.type != http_websocket.WSMsgType.TEXT:
                continue

            coro = self._handle_ws_message(
                ws_msg=ws_msg,
                ws_connect=ws_connect,
                context={
                    'http_request': http_request,
                    'ws_connect': ws_connect,
                    'ws_rpc_client': ws_rpc_client,
                },
            )

            task = asyncio.create_task(coro)

            # To avoid a task disappearing mid execution:
            self._background_tasks.add(task)
            task.add_done_callback(self._background_tasks.discard)

        return ws_connect

    async def _handle_ws_message(self,
                                 ws_msg: web_ws.WSMessage, *,
                                 ws_connect: web_ws.WebSocketResponse,
                                 context: dict) -> None:
        json_response: typing.Optional[typing.Union[typing.Mapping, typing.Sequence[typing.Mapping]]]

        try:
            input_data = json.loads(ws_msg.data)
        except json.JSONDecodeError as e:
            response = protocol.JsonRpcResponse(error=errors.ParseError(utils.get_exc_message(e)))
            json_response = response.dump()
        else:
            json_response = await self._process_input_data(input_data, context=context)

        if json_response is None:
            return

        if ws_connect.closed:
            raise errors.ServerError('WS is closed.')

        await ws_connect.send_str(self.json_serialize(json_response))