Add ws_connect to ClientSession (cleaned up version of pull request #371 - figuring out the right way to do this)

aiohttp/client.py

@@ -1,15 +1,21 @@
 """HTTP Client for asyncio."""
 
 import asyncio
-import http.cookies
-import urllib.parse
-import warnings
+import base64
+import hashlib
+import os
 import sys
 import traceback
+import warnings
+import http.cookies
+import urllib.parse
 
 import aiohttp
 from .client_reqrep import ClientRequest, ClientResponse
+from .errors import WSServerHandshakeError
 from .multidict import MultiDictProxy, MultiDict, CIMultiDict
+from .websocket import WS_KEY, WebSocketParser, WebSocketWriter
+from .websocket_client import ClientWebSocketResponse
 from . import hdrs
 
 
@@ -26,7 +32,8 @@ class ClientSession:
 
     def __init__(self, *, connector=None, loop=None,
                  request_class=ClientRequest, response_class=ClientResponse,
-                 cookies=None, headers=None, auth=None):
+                 ws_response_class=ClientWebSocketResponse, cookies=None,
+                 headers=None, auth=None):
         if loop is None:
             loop = asyncio.get_event_loop()
         self._loop = loop
@@ -58,6 +65,7 @@ def __init__(self, *, connector=None, loop=None,
 
         self._request_class = request_class
         self._response_class = response_class
+        self._ws_response_class = ws_response_class
 
     if PY_34:
         def __del__(self):
@@ -167,6 +175,70 @@ def request(self, method, url, *,
 
         return resp
 
+    @asyncio.coroutine
+    def ws_connect(self, url, *,
+                   protocols=(),
+                   timeout=10.0,
+                   autoclose=True,
+                   autoping=True):
+        """Initiate websocket connection."""
+
+        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 self.request('get', url, headers=headers,
+                                       read_until_eof=False)
+
+        # 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)
+
+        return self._ws_response_class(reader,
+                                       writer,
+                                       protocol,
+                                       resp,
+                                       timeout,
+                                       autoclose,
+                                       autoping,
+                                       self._loop)
+
     def _update_cookies(self, cookies):
         """Update shared cookies."""
         if isinstance(cookies, dict):

aiohttp/websocket_client.py

@@ -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,34 @@ 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,
+               ws_response_class=None, autoclose=True, autoping=True,
+               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)
+
+    kwargs = {}
+
+    if ws_response_class is not None:
+        kwargs['ws_response_class'] = ws_response_class
+
+    session = aiohttp.ClientSession(loop=loop, connector=connector, **kwargs)
+
+    try:
+        resp = yield from session.ws_connect(
+            url,
+            protocols=protocols,
+            timeout=timeout,
+            autoclose=autoclose,
+            autoping=autoping)
+        return resp
+
+    finally:
+        session.detach()
 
 
 class ClientWebSocketResponse:

docs/client_reference.rst

@@ -30,8 +30,8 @@ Usage example::
 
 
 .. class:: ClientSession(*, connector=None, loop=None, request_class=None,\
-                          response_class=None, cookies=None, headers=None,\
-                          auth=None)
+                          response_class=None, ws_response_class=None,\
+                          cookies=None, headers=None, auth=None)
 
    The class for creating client sessions and making requests.
 
@@ -51,6 +51,9 @@ Usage example::
 
    :param response_class: Custom Response class implementation (optional)
 
+   :param ws_response_class: Custom WebSocketResponse class implementation
+                             (optional)
+
    :param dict cookies: Cookies to send with the request (optional)
 
    :param dict headers: HTTP Headers to send with
@@ -223,6 +226,25 @@ 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\
+                                   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 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
+
+
    .. method:: close()
 
       Close underlying connector.

docs/client_websockets.rst

@@ -46,25 +46,30 @@ 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,\
+                            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 WebSocketResponse 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.