Implement optional nursery pattern for client (#31)

In addition to open_websocket and open_websocket_url, which are
both context managers that do not have a nursery argument, I have
added connect_websocket() and connect_websocket_url(), which take
a nursery argument and are not context managers.

tests/test_connection.py

@@ -1,8 +1,8 @@
 import logging
 
 import pytest
-from trio_websocket import ConnectionClosed, open_websocket, \
-    open_websocket_url, WebSocketServer
+from trio_websocket import ConnectionClosed, connect_websocket, \
+    connect_websocket_url, open_websocket, open_websocket_url, WebSocketServer
 import trio
 
 
@@ -34,24 +34,36 @@ async def echo_conn(echo_server):
         yield conn
 
 
+async def test_client_open(echo_server):
+    async with open_websocket(HOST, echo_server.port, RESOURCE, use_ssl=False) \
+        as conn:
+        assert conn.closed is None
+
+
 async def test_client_open_url(echo_server):
     url = 'ws://{}:{}{}?foo=bar'.format(HOST, echo_server.port, RESOURCE)
     async with open_websocket_url(url) as conn:
         assert conn.path == RESOURCE + '?foo=bar'
 
 
-async def test_client_open_in_nursery(echo_server, nursery):
-    url = 'ws://{}:{}{}'.format(HOST, echo_server.port, RESOURCE)
-    async with open_websocket_url(url, nursery=nursery) as conn:
-        assert len(nursery.child_tasks) == 1
-
-
 async def test_client_open_invalid_url(echo_server):
     with pytest.raises(ValueError):
         async with open_websocket_url('http://foo.com/bar') as conn:
             pass
 
 
+async def test_client_connect(echo_server, nursery):
+    conn = await connect_websocket(nursery, HOST, echo_server.port, RESOURCE,
+        use_ssl=False)
+    assert conn.closed is None
+
+
+async def test_client_connect_url(echo_server, nursery):
+    url = 'ws://{}:{}{}'.format(HOST, echo_server.port, RESOURCE)
+    conn = await connect_websocket_url(nursery, url)
+    assert conn.closed is None
+
+
 async def test_client_send_and_receive(echo_conn):
     async with echo_conn:
         await echo_conn.send_message('This is a test message.')

trio_websocket/__init__.py

@@ -22,21 +22,46 @@
 
 @asynccontextmanager
 @async_generator
-async def open_websocket(host, port, resource, use_ssl, nursery=None):
+async def open_websocket(host, port, resource, use_ssl):
     '''
     Open a WebSocket client connection to a host.
 
-    This function is an async context manager that automatically connects and
-    disconnects. It yields a `WebSocketConnection` instance.
+    This function is an async context manager that connects before entering the
+    context manager and disconnects after leaving. It yields a
+    `WebSocketConnection` instance.
+
+    ``use_ssl`` can be an ``SSLContext`` object, or if it's ``True``, then a
+    default SSL context will be created. If it is ``False``, then SSL will not
+    be used at all.
 
     :param str host: the host to connect to
     :param int port: the port to connect to
     :param str resource: the resource a.k.a. path
     :param use_ssl: a bool or SSLContext
-    :param nursery: optional Trio nursery to run background tasks in, if not
-        provided, then a new nursery will be created
     '''
+    async with trio.open_nursery() as new_nursery:
+        connection = await connect_websocket(new_nursery, host, port, resource,
+            use_ssl)
+        async with connection:
+            await yield_(connection)
+
+
+async def connect_websocket(nursery, host, port, resource, use_ssl):
+    '''
+    Return a WebSocket client connection to a host.
 
+    Most users should use ``open_websocket(…)`` instead of this function. This
+    function is not an async context manager, and it requires a nursery argument
+    for the connection's background task[s]. The caller is responsible for
+    closing the connection.
+
+    :param nursery: a Trio nursery to run background tasks in
+    :param str host: the host to connect to
+    :param int port: the port to connect to
+    :param str resource: the resource a.k.a. path
+    :param use_ssl: a bool or SSLContext
+    :rtype: WebSocketConnection
+    '''
     if use_ssl == True:
         ssl_context = ssl.create_default_context()
     elif use_ssl == False:
@@ -60,44 +85,72 @@ async def open_websocket(host, port, resource, use_ssl, nursery=None):
     wsproto = wsconnection.WSConnection(wsconnection.CLIENT,
         host=host_header, resource=resource)
     connection = WebSocketConnection(stream, wsproto, path=resource)
+    nursery.start_soon(connection._reader_task)
+    await connection._open_handshake.wait()
+    return connection
 
-    async def connect_in_nursery(connection, nursery):
-        nursery.start_soon(connection._reader_task)
-        await connection._open_handshake.wait()
-        return connection
 
-    if nursery is None:
-        async with trio.open_nursery() as new_nursery:
-            connection = await connect_in_nursery(connection, new_nursery)
-            async with connection:
-                await yield_(connection)
-    else:
-        connection = await connect_in_nursery(connection, nursery)
-        async with connection:
-            await yield_(connection)
+def open_websocket_url(url, ssl_context=None):
+    '''
+    Open a WebSocket client connection to a URL.
+
+    This function is an async context manager that connects when entering the
+    context manager and disconnects when exiting. It yields a
+    `WebSocketConnection` instance.
 
+    If ``ssl_context`` is ``None`` and the URL scheme is ``wss:``, then a
+    default SSL context will be created. It is an error to pass an SSL context
+    for ``ws:`` URLs.
 
-def open_websocket_url(url, ssl_context=None, nursery=None):
+    :param str url: a WebSocket URL
+    :param ssl_context: optional ``SSLContext`` used for ``wss:`` URLs
     '''
-    Open a WebSocket client connection to a URL.
+    host, port, resource, ssl_context = _url_to_host(url, ssl_context)
+    return open_websocket(host, port, resource, ssl_context)
+
+
+async def connect_websocket_url(nursery, url, ssl_context=None):
+    '''
+    Return a WebSocket client connection to a URL.
+
+    Most users should use ``open_websocket_url(…)`` instead of this function.
+    This function is not an async context manager, and it requires a nursery
+    argument for the connection's background task[s].
 
-    This function is an async context manager that automatically connects and
-    disconnects. It yields a `WebSocketConnection` instance.
+    If ``ssl_context`` is ``None`` and the URL scheme is ``wss:``, then a
+    default SSL context will be created. It is an error to pass an SSL context
+    for ``ws:`` URLs.
 
     :param str url: a WebSocket URL
-    :param ssl_context: optional SSLContext, only used for wss: URL scheme
-    :param nursery: optional Trio nursery to run background tasks in, if not
-        provided, then a new nursery will be created
+    :param ssl_context: optional ``SSLContext`` used for ``wss:`` URLs
+    :param nursery: a Trio nursery to run background tasks in
+    :rtype: WebSocketConnection
+    '''
+    host, port, resource, ssl_context = _url_to_host(url, ssl_context)
+    return await connect_websocket(nursery, host, port, resource, ssl_context)
+
+
+def _url_to_host(url, ssl_context):
+    '''
+    Convert a WebSocket URL to a (host,port,resource) tuple.
+
+    The returned ``ssl_context`` is either the same object that was passed in,
+    or if ``ssl_context`` is None, then a bool indicating if a default SSL
+    context needs to be created.
+
+    :param str url: a WebSocket URL
+    :param ssl_context: ``SSLContext`` or ``None``
+    :return: tuple of ``(host, port, resource, ssl_context)``
     '''
     url = URL(url)
     if url.scheme not in ('ws', 'wss'):
         raise ValueError('WebSocket URL scheme must be "ws:" or "wss:"')
     if ssl_context is None:
-        use_ssl = url.scheme == 'wss'
-    else:
-        use_ssl = ssl_context
+        ssl_context = url.scheme == 'wss'
+    elif url.scheme == 'ws':
+        raise ValueError('SSL context must be None for ws: URL scheme')
     resource = '{}?{}'.format(url.path, url.query_string)
-    return open_websocket(url.host, url.port, resource, use_ssl, nursery)
+    return url.host, url.port, resource, ssl_context
 
 
 class ConnectionClosed(Exception):