bazarr/libs/simple_websocket/asgi.py

51 lines
1.8 KiB
Python

from .errors import ConnectionClosed # pragma: no cover
class WebSocketASGI: # pragma: no cover
def __init__(self, scope, receive, send, subprotocols=None):
self._scope = scope
self._receive = receive
self._send = send
self.subprotocols = subprotocols or []
self.subprotocol = None
self.connected = False
@classmethod
async def accept(cls, scope, receive, send, subprotocols=None):
ws = WebSocketASGI(scope, receive, send, subprotocols=subprotocols)
await ws._accept()
return ws
async def _accept(self):
connect = await self._receive()
if connect['type'] != 'websocket.connect':
raise ValueError('Expected websocket.connect')
for subprotocol in self._scope['subprotocols']:
if subprotocol in self.subprotocols:
self.subprotocol = subprotocol
break
await self._send({'type': 'websocket.accept',
'subprotocol': self.subprotocol})
async def receive(self):
message = await self._receive()
if message['type'] == 'websocket.disconnect':
raise ConnectionClosed()
elif message['type'] != 'websocket.receive':
raise OSError(32, 'Websocket message type not supported')
return message.get('text', message.get('bytes'))
async def send(self, data):
if isinstance(data, str):
await self._send({'type': 'websocket.send', 'text': data})
else:
await self._send({'type': 'websocket.send', 'bytes': data})
async def close(self):
if not self.connected:
self.conncted = False
try:
await self._send({'type': 'websocket.close'})
except Exception:
pass