codekingpro/portable-devtools

1from .errors import ConnectionClosed  # pragma: no cover2 3 4class WebSocketASGI:  # pragma: no cover5    def __init__(self, scope, receive, send, subprotocols=None):6        self._scope = scope7        self._receive = receive8        self._send = send9        self.subprotocols = subprotocols or []10        self.subprotocol = None11        self.connected = False12 13    @classmethod14    async def accept(cls, scope, receive, send, subprotocols=None):15        ws = WebSocketASGI(scope, receive, send, subprotocols=subprotocols)16        await ws._accept()17        return ws18 19    async def _accept(self):20        connect = await self._receive()21        if connect['type'] != 'websocket.connect':22            raise ValueError('Expected websocket.connect')23        for subprotocol in self._scope['subprotocols']:24            if subprotocol in self.subprotocols:25                self.subprotocol = subprotocol26                break27        await self._send({'type': 'websocket.accept',28                         'subprotocol': self.subprotocol})29 30    async def receive(self):31        message = await self._receive()32        if message['type'] == 'websocket.disconnect':33            raise ConnectionClosed()34        elif message['type'] != 'websocket.receive':35            raise OSError(32, 'Websocket message type not supported')36        return message.get('text', message.get('bytes'))37 38    async def send(self, data):39        if isinstance(data, str):40            await self._send({'type': 'websocket.send', 'text': data})41        else:42            await self._send({'type': 'websocket.send', 'bytes': data})43 44    async def close(self):45        if not self.connected:46            self.conncted = False47            try:48                await self._send({'type': 'websocket.close'})49            except Exception:50                pass51