You cannot select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
Enso-Bot/venv/Lib/site-packages/aiohttp/client_ws.py

194 lines
6.1 KiB
Python

"""WebSocket client for asyncio."""
import asyncio
import json
import sys
from ._ws_impl import CLOSED_MESSAGE, WebSocketError, WSMessage, WSMsgType
PY_35 = sys.version_info >= (3, 5)
PY_352 = sys.version_info >= (3, 5, 2)
class ClientWebSocketResponse:
def __init__(self, reader, writer, protocol,
response, timeout, autoclose, autoping, loop):
self._response = response
self._conn = response.connection
self._writer = writer
self._reader = reader
self._protocol = protocol
self._closed = False
self._closing = False
self._close_code = None
self._timeout = timeout
self._autoclose = autoclose
self._autoping = autoping
self._loop = loop
self._waiting = False
self._exception = None
@property
def closed(self):
return self._closed
@property
def close_code(self):
return self._close_code
@property
def protocol(self):
return self._protocol
def exception(self):
return self._exception
def ping(self, message='b'):
if self._closed:
raise RuntimeError('websocket connection is closed')
self._writer.ping(message)
def pong(self, message='b'):
if self._closed:
raise RuntimeError('websocket connection is closed')
self._writer.pong(message)
def send_str(self, data):
if self._closed:
raise RuntimeError('websocket connection is closed')
if not isinstance(data, str):
raise TypeError('data argument must be str (%r)' % type(data))
self._writer.send(data, binary=False)
def send_bytes(self, data):
if self._closed:
raise RuntimeError('websocket connection is closed')
if not isinstance(data, (bytes, bytearray, memoryview)):
raise TypeError('data argument must be byte-ish (%r)' %
type(data))
self._writer.send(data, binary=True)
def send_json(self, data, *, dumps=json.dumps):
self.send_str(dumps(data))
@asyncio.coroutine
def close(self, *, code=1000, message=b''):
if not self._closed:
self._closed = True
try:
self._writer.close(code, message)
except asyncio.CancelledError:
self._close_code = 1006
self._response.close()
raise
except Exception as exc:
self._close_code = 1006
self._exception = exc
self._response.close()
return True
if self._closing:
self._response.close()
return True
while True:
try:
msg = yield from asyncio.wait_for(
self._reader.read(), self._timeout, loop=self._loop)
except asyncio.CancelledError:
self._close_code = 1006
self._response.close()
raise
except Exception as exc:
self._close_code = 1006
self._exception = exc
self._response.close()
return True
if msg.type == WSMsgType.CLOSE:
self._close_code = msg.data
self._response.close()
return True
else:
return False
@asyncio.coroutine
def receive(self):
if self._waiting:
raise RuntimeError('Concurrent call to receive() is not allowed')
self._waiting = True
try:
while True:
if self._closed:
return CLOSED_MESSAGE
try:
msg = yield from self._reader.read()
except (asyncio.CancelledError, asyncio.TimeoutError):
raise
except WebSocketError as exc:
self._close_code = exc.code
yield from self.close(code=exc.code)
return WSMessage(WSMsgType.ERROR, exc, None)
except Exception as exc:
self._exception = exc
self._closing = True
self._close_code = 1006
yield from self.close()
return WSMessage(WSMsgType.ERROR, exc, None)
if msg.type == WSMsgType.CLOSE:
self._closing = True
self._close_code = msg.data
if not self._closed and self._autoclose:
yield from self.close()
return msg
if msg.type == WSMsgType.PING and self._autoping:
self.pong(msg.data)
elif msg.type == WSMsgType.PONG and self._autoping:
continue
else:
return msg
finally:
self._waiting = False
@asyncio.coroutine
def receive_str(self):
msg = yield from self.receive()
if msg.type != WSMsgType.TEXT:
raise TypeError(
"Received message {}:{!r} is not str".format(msg.type,
msg.data))
return msg.data
@asyncio.coroutine
def receive_bytes(self):
msg = yield from self.receive()
if msg.type != WSMsgType.BINARY:
raise TypeError(
"Received message {}:{!r} is not bytes".format(msg.type,
msg.data))
return msg.data
@asyncio.coroutine
def receive_json(self, *, loads=json.loads):
data = yield from self.receive_str()
return loads(data)
if PY_35:
def __aiter__(self):
return self
if not PY_352: # pragma: no cover
__aiter__ = asyncio.coroutine(__aiter__)
@asyncio.coroutine
def __anext__(self):
msg = yield from self.receive()
if msg.type == WSMsgType.CLOSE:
raise StopAsyncIteration # NOQA
return msg