New issue
Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.
By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.
Already on GitHub? Sign in to your account
Add ws_connect to ClientSession (cleaned up version of pull request #371 - figuring out the right way to do this) #374
Changes from 1 commit
2545cef
29fee1d
ae43d46
8188ebc
a953607
File filter
Filter by extension
Conversations
Jump to
Diff view
Diff view
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -1,14 +1,10 @@ | ||
"""WebSocket client for asyncio.""" | ||
|
||
import asyncio | ||
import base64 | ||
import hashlib | ||
import os | ||
|
||
from aiohttp import client, hdrs | ||
from .errors import WSServerHandshakeError | ||
from .websocket import WS_KEY, Message | ||
from .websocket import WebSocketParser, WebSocketWriter, WebSocketError | ||
|
||
import aiohttp | ||
from .websocket import Message | ||
from .websocket import WebSocketError | ||
from .websocket import MSG_BINARY, MSG_TEXT, MSG_CLOSE, MSG_PING, MSG_PONG | ||
|
||
__all__ = ('ws_connect', 'MsgType') | ||
|
@@ -33,65 +29,30 @@ class MsgType(IntEnum): | |
closedMessage = Message(MsgType.closed, None, None) | ||
|
||
|
||
@asyncio.coroutine | ||
def ws_connect(url, protocols=(), timeout=10.0, connector=None, | ||
response_class=None, autoclose=True, autoping=True, loop=None): | ||
"""Initiate websocket connection.""" | ||
def ws_connect(url, *, protocols=(), timeout=10.0, connector=None, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Please add |
||
ws_response_class=None, autoclose=True, autoping=True, | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. Don't hesitate to use There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I'm unsure about this. Would we still check that ws_connect("ws...", response_class=None) then _self.ws_response_class on the There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. I guess that if they change the param we have to assume they know what they are doing anyway. |
||
loop=None): | ||
|
||
if loop is None: | ||
loop = asyncio.get_event_loop() | ||
|
||
sec_key = base64.b64encode(os.urandom(16)) | ||
|
||
headers = { | ||
hdrs.UPGRADE: hdrs.WEBSOCKET, | ||
hdrs.CONNECTION: hdrs.UPGRADE, | ||
hdrs.SEC_WEBSOCKET_VERSION: '13', | ||
hdrs.SEC_WEBSOCKET_KEY: sec_key.decode(), | ||
} | ||
if protocols: | ||
headers[hdrs.SEC_WEBSOCKET_PROTOCOL] = ','.join(protocols) | ||
|
||
# send request | ||
resp = yield from client.request( | ||
'get', url, headers=headers, | ||
read_until_eof=False, | ||
connector=connector, loop=loop) | ||
|
||
# check handshake | ||
if resp.status != 101: | ||
raise WSServerHandshakeError('Invalid response status') | ||
|
||
if resp.headers.get(hdrs.UPGRADE, '').lower() != 'websocket': | ||
raise WSServerHandshakeError('Invalid upgrade header') | ||
|
||
if resp.headers.get(hdrs.CONNECTION, '').lower() != 'upgrade': | ||
raise WSServerHandshakeError('Invalid connection header') | ||
|
||
# key calculation | ||
key = resp.headers.get(hdrs.SEC_WEBSOCKET_ACCEPT, '') | ||
match = base64.b64encode(hashlib.sha1(sec_key + WS_KEY).digest()).decode() | ||
if key != match: | ||
raise WSServerHandshakeError('Invalid challenge response') | ||
|
||
# websocket protocol | ||
protocol = None | ||
if protocols and hdrs.SEC_WEBSOCKET_PROTOCOL in resp.headers: | ||
resp_protocols = [proto.strip() for proto in | ||
resp.headers[hdrs.SEC_WEBSOCKET_PROTOCOL].split(',')] | ||
|
||
for proto in resp_protocols: | ||
if proto in protocols: | ||
protocol = proto | ||
break | ||
|
||
reader = resp.connection.reader.set_parser(WebSocketParser) | ||
writer = WebSocketWriter(resp.connection.writer, use_mask=True) | ||
|
||
if response_class is None: | ||
response_class = ClientWebSocketResponse | ||
|
||
return response_class( | ||
reader, writer, protocol, resp, timeout, autoclose, autoping, loop) | ||
asyncio.get_event_loop() | ||
|
||
if connector is None: | ||
connector = aiohttp.TCPConnector(loop=loop, force_close=True) | ||
|
||
session = aiohttp.ClientSession(loop=loop, connector=connector) | ||
|
||
try: | ||
resp = yield from session.ws_connect( | ||
url, | ||
protocols=protocols, | ||
timeout=timeout, | ||
ws_response_class=ws_response_class, | ||
autoclose=autoclose, | ||
autoping=autoping) | ||
return resp | ||
|
||
finally: | ||
session.detach() | ||
|
||
|
||
class ClientWebSocketResponse: | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -223,6 +223,28 @@ Usage example:: | |
:param data: Dictionary, bytes, or file-like object to | ||
send in the body of the request (optional) | ||
|
||
|
||
.. coroutinemethod:: ws_connect(url, *, protocols=(), timeout=10.0\ | ||
ws_response_class=None, autoclose=True,\ | ||
autoping=True) | ||
|
||
Create a websocket connection. Returns a :class:`ClientWebSocketResponse` object. | ||
|
||
:param str url: Websocket server url | ||
|
||
:param tuple protocols: Websocket protocols | ||
|
||
:param float timeout: Timeout for websocket read. 10 seconds by default | ||
|
||
:param ws_response_class: (optional) Custom Response class implementation | ||
|
||
:param bool autoclose: Automatically close websocket connection on close | ||
message from server. If `autoclose` is False | ||
them close procedure has to be handled manually | ||
|
||
:param bool autoping: automatically send `pong` on `ping` message from server | ||
|
||
|
||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more.
|
||
.. method:: close() | ||
|
||
Close underlying connector. | ||
|
Original file line number | Diff line number | Diff line change |
---|---|---|
|
@@ -46,25 +46,29 @@ ClientWebSocketResponse | |
To connect to a websocket server you have to use the `aiohttp.ws_connect()` function, | ||
do not create an instance of class :class:`ClientWebSocketResponse` manually. | ||
|
||
.. py:function:: ws_connect(url, protocols=(), connector=None, response_class=None, autoclose=True, autoping=True, loop=None) | ||
.. py:function:: ws_connect(url, *, protocols=(), timeout=10.0,\ | ||
connector=None, ws_response_class=None,\ | ||
There was a problem hiding this comment. Choose a reason for hiding this commentThe reason will be displayed to describe this comment to others. Learn more. ws_response_class=ClientWebSocketResponse |
||
autoclose=True, autoping=True, loop=None) | ||
|
||
This function creates a websocket connection, checks the response and | ||
returns a :class:`ClientWebSocketResponse` object. In case of failure | ||
it may raise a :exc:`~aiohttp.errors.WSServerHandshakeError` exception. | ||
|
||
:param str url: websocket server url | ||
:param str url: Websocket server url | ||
|
||
:param tuple protocols: websocket protocols | ||
:param tuple protocols: Websocket protocols | ||
|
||
:param float timeout: Timeout for websocket read. 10 seconds by default | ||
|
||
:param obj connector: object :class:`TCPConnector` | ||
|
||
:param response_class: (optional) Custom Response class implementation. | ||
:param ws_response_class: (optional) Custom Response class implementation. | ||
|
||
:param bool autoclose: automatically close websocket connection | ||
on close message from server. if `autoclose` is | ||
:param bool autoclose: Automatically close websocket connection | ||
on close message from server. If `autoclose` is | ||
False them close procedure has to be handled manually | ||
|
||
:param bool autoping: automatically send `pong` on `ping` message from server | ||
:param bool autoping: Automatically send `pong` on `ping` message from server | ||
|
||
:param loop: :ref:`event loop<asyncio-event-loop>` used | ||
for processing HTTP requests. | ||
|
There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
The only possible other change I would make to this patch would be to move ws_response_class to the constructor function to be consistent with
ClientSession.request
method.There was a problem hiding this comment.
Choose a reason for hiding this comment
The reason will be displayed to describe this comment to others. Learn more.
Agree. Please do.