asyncio: Add streams.start_server(), by Gustavo Carneiro.

1 files changed, 51 insertions, 2 deletions

diff --git a/Lib/asyncio/streams.py b/Lib/asyncio/streams.py
index e995368..331d28d 100644
--- a/Lib/asyncio/streams.py
+++ b/Lib/asyncio/streams.py
@@ -1,6 +1,8 @@
 """Stream-related things."""
 
-__all__ = ['StreamReader', 'StreamReaderProtocol', 'open_connection']
+__all__ = ['StreamReader', 'StreamReaderProtocol',
+           'open_connection', 'start_server',
+           ]
 
 import collections
 
@@ -43,6 +45,42 @@ def open_connection(host=None, port=None, *,
     return reader, writer
 
 
+@tasks.coroutine
+def start_server(client_connected_cb, host=None, port=None, *,
+                 loop=None, limit=_DEFAULT_LIMIT, **kwds):
+    """Start a socket server, call back for each client connected.
+
+    The first parameter, `client_connected_cb`, takes two parameters:
+    client_reader, client_writer.  client_reader is a StreamReader
+    object, while client_writer is a StreamWriter object.  This
+    parameter can either be a plain callback function or a coroutine;
+    if it is a coroutine, it will be automatically converted into a
+    Task.
+
+    The rest of the arguments are all the usual arguments to
+    loop.create_server() except protocol_factory; most common are
+    positional host and port, with various optional keyword arguments
+    following.  The return value is the same as loop.create_server().
+
+    Additional optional keyword arguments are loop (to set the event loop
+    instance to use) and limit (to set the buffer limit passed to the
+    StreamReader).
+
+    The return value is the same as loop.create_server(), i.e. a
+    Server object which can be used to stop the service.
+    """
+    if loop is None:
+        loop = events.get_event_loop()
+
+    def factory():
+        reader = StreamReader(limit=limit, loop=loop)
+        protocol = StreamReaderProtocol(reader, client_connected_cb,
+                                        loop=loop)
+        return protocol
+
+    return (yield from loop.create_server(factory, host, port, **kwds))
+
+
 class StreamReaderProtocol(protocols.Protocol):
     """Trivial helper class to adapt between Protocol and StreamReader.
 
@@ -52,13 +90,24 @@ class StreamReaderProtocol(protocols.Protocol):
     call inappropriate methods of the protocol.)
     """
 
-    def __init__(self, stream_reader):
+    def __init__(self, stream_reader, client_connected_cb=None, loop=None):
         self._stream_reader = stream_reader
+        self._stream_writer = None
         self._drain_waiter = None
         self._paused = False
+        self._client_connected_cb = client_connected_cb
+        self._loop = loop  # May be None; we may never need it.
 
     def connection_made(self, transport):
         self._stream_reader.set_transport(transport)
+        if self._client_connected_cb is not None:
+            self._stream_writer = StreamWriter(transport, self,
+                                               self._stream_reader,
+                                               self._loop)
+            res = self._client_connected_cb(self._stream_reader,
+                                            self._stream_writer)
+            if tasks.iscoroutine(res):
+                tasks.Task(res, loop=self._loop)
 
     def connection_lost(self, exc):
         if exc is None: