Make serve an async context manager.

Fix python-websockets#86.

Author: aaugustin

docs/changelog.rst

@@ -6,6 +6,9 @@ Changelog
 
 *In development*
 
+* :func:`~websockets.server.serve` can be used as an asynchronous context
+  manager on Python ≥ 3.5.
+
 * Added support rejecting incoming connections by customizing
   :meth:`~websockets.server.WebSocketServerProtocol.get_response_status()`.
 

docs/deployment.rst

@@ -91,11 +91,12 @@ Graceful shutdown
 -----------------
 
 You may want to close connections gracefully when shutting down the server,
-perhaps after executing some cleanup logic.
+perhaps after executing some cleanup logic. There are two ways to achieve this
+with the object returned by :func:`~websockets.server.serve`:
 
-The proper way to do this is to call the ``close()`` method of the object
-returned by :func:`~websockets.server.serve`, then wait for ``wait_closed()``
-to complete.
+- using it as a asynchronous context manager, or
+- calling its ``close()`` method, then waiting for its ``wait_closed()``
+  method to complete.
 
 Tasks that handle connections will be cancelled, in the sense that
 :meth:`~websockets.protocol.WebSocketCommonProtocol.recv` raises
@@ -107,6 +108,11 @@ Here's a full example (Unix-only):
 
 .. literalinclude:: ../example/shutdown.py
 
+``async``, ``await``, and asynchronous context managers aren't available in
+Python < 3.5. Here's the equivalent for older Python versions:
+
+.. literalinclude:: ../example/oldshutdown.py
+
 It's more difficult to achieve the same effect on Windows. Some third-party
 projects try to help with this problem.
 

example/oldshutdown.py

@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+
+import asyncio
+import signal
+import websockets
+
+async def echo(websocket, path):
+    while True:
+        try:
+            msg = await websocket.recv()
+        except websockets.ConnectionClosed:
+            pass
+        else:
+            await websocket.send(msg)
+
+loop = asyncio.get_event_loop()
+
+# Create the server.
+start_server = websockets.serve(echo, 'localhost', 8765)
+server = loop.run_until_complete(start_server)
+
+# Run the server until SIGTERM.
+stop = asyncio.Future()
+loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
+loop.run_until_complete(stop)
+
+# Shut down the server.
+server.close()
+loop.run_until_complete(server.wait_closed())

example/shutdown.py

@@ -13,17 +13,15 @@ async def echo(websocket, path):
         else:
             await websocket.send(msg)
 
-loop = asyncio.get_event_loop()
+async def echo_server(stop):
+    async with websockets.serve(echo, 'localhost', 8765):
+        await stop
 
-# Create the server.
-start_server = websockets.serve(echo, 'localhost', 8765)
-server = loop.run_until_complete(start_server)
+loop = asyncio.get_event_loop()
 
-# Run the server until SIGTERM.
-stop_server = asyncio.Future()
-loop.add_signal_handler(signal.SIGTERM, stop_server.set_result, None)
-loop.run_until_complete(stop_server)
+# The stop condition is set when receiving SIGTERM.
+stop = asyncio.Future()
+loop.add_signal_handler(signal.SIGTERM, stop.set_result, None)
 
-# Shut down the server.
-server.close()
-loop.run_until_complete(server.wait_closed())
+# Run the server until the stop condition is met.
+loop.run_until_complete(echo_server(stop))

websockets/client.py

@@ -147,7 +147,7 @@ def connect(uri, *,
     send and receive messages.
 
     :func:`connect` is a wrapper around the event loop's
-    :meth:`~asyncio.BaseEventLoop.create_connection` method. Extra keyword
+    :meth:`~asyncio.BaseEventLoop.create_connection` method. Unknown keyword
     arguments are passed to :meth:`~asyncio.BaseEventLoop.create_connection`.
 
     For example, you can set the ``ssl`` keyword argument to a

websockets/py35/client_server.py

@@ -17,17 +17,27 @@ def setUp(self):
     def tearDown(self):
         self.loop.close()
 
-    def test_basic(self):
+    def test_client(self):
         server = serve(handler, 'localhost', 8642)
         self.server = self.loop.run_until_complete(server)
 
-        async def basic():
+        async def run_client():
             async with connect('ws://localhost:8642/') as client:
                 await client.send("Hello!")
                 reply = await client.recv()
                 self.assertEqual(reply, "Hello!")
 
-        self.loop.run_until_complete(basic())
+        self.loop.run_until_complete(run_client())
 
         self.server.close()
         self.loop.run_until_complete(self.server.wait_closed())
+
+    def test_server(self):
+        async def run_server():
+            async with serve(handler, 'localhost', 8642):
+                client = await connect('ws://localhost:8642/')
+                await client.send("Hello!")
+                reply = await client.recv()
+                self.assertEqual(reply, "Hello!")
+
+        self.loop.run_until_complete(run_server())

websockets/py35/server.py

@@ -0,0 +1,22 @@
+class Serve:
+    """
+    This class wraps :func:`~websockets.server.serve` on Python ≥ 3.5.
+
+    This allows using it as an asynchronous context manager.
+
+    """
+    def __init__(self, *args, **kwargs):
+        self.server = self.__class__.__wrapped__(*args, **kwargs)
+
+    async def __aenter__(self):
+        self.server = await self
+        return self.server
+
+    async def __aexit__(self, exc_type, exc_value, traceback):
+        self.server.close()
+        await self.server.wait_closed()
+
+    def __await__(self):
+        return (yield from self.server)
+
+    __iter__ = __await__

websockets/server.py

@@ -407,7 +407,7 @@ def serve(ws_handler, host=None, port=None, *,
 
     :func:`serve` is a wrapper around the event loop's
     :meth:`~asyncio.BaseEventLoop.create_server` method. ``host``, ``port`` as
-    well as extra keyword arguments are passed to
+    well as unknown keyword arguments are passed to
     :meth:`~asyncio.BaseEventLoop.create_server`.
 
     For example, you can set the ``ssl`` keyword argument to a
@@ -441,6 +441,9 @@ def serve(ws_handler, host=None, port=None, *,
         logger.setLevel(logging.ERROR)
         logger.addHandler(logging.StreamHandler())
 
+    On Python 3.5, :func:`serve` can be used as a asynchronous context
+    manager. In that case, the server is shut down when exiting the context.
+
     """
     if loop is None:
         loop = asyncio.get_event_loop()
@@ -462,3 +465,14 @@ def serve(ws_handler, host=None, port=None, *,
     ws_server.wrap(server)
 
     return ws_server
+
+
+try:
+    from .py35.server import Serve
+except (SyntaxError, ImportError):                          # pragma: no cover
+    pass
+else:
+    Serve.__wrapped__ = serve
+    # Copy over docstring to support building documentation on Python 3.5.
+    Serve.__doc__ = serve.__doc__
+    serve = Serve