Repository: aaugustin/websockets (python)

Issue:
Support async iterators for fragmenting outgoing frames
This is a follow-up to #258.

`collections.abc.AsyncIterable` is new in Python 3.5. This needs dropping support for Python 3.4 first to keep things simple.

Also it requires preventing concurrent calls to `send()`, else another call could insert a data frame while we're sending fragments.

Pull request:
Send fragmented messages from async iterators.
Fix #477.

Patch:
diff --git a/src/websockets/protocol.py b/src/websockets/protocol.py
index 981c097..8dacbf4 100644
--- a/src/websockets/protocol.py
+++ b/src/websockets/protocol.py
@@ -432,10 +432,11 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
         object (:class:`bytes`, :class:`bytearray`, or :class:`memoryview`)
         as a binary frame.
 
-        It also accepts an iterable of strings or bytes-like objects. Each
-        item is treated as a message fragment and sent in its own frame. All
-        items must be of the same type, or else :meth:`send` will raise a
-        :exc:`TypeError` and the connection will be closed.
+        It also accepts an iterable or an asynchronous iterator of strings or
+        bytes-like objects. Each item is treated as a message fragment and
+        sent in its own frame. All items must be of the same type, or else
+        :meth:`send` will raise a :exc:`TypeError` and the connection will be
+        closed.
 
         It raises a :exc:`TypeError` for other inputs.
 
@@ -482,7 +483,31 @@ class WebSocketCommonProtocol(asyncio.StreamReaderProtocol):
 
         # Fragmented message -- asynchronous iterator
 
-        # To be implemented after dropping support for Python 3.4.
+        elif isinstance(data, collections.abc.AsyncIterable):
+            # aiter_data = aiter(data) without aiter
+            aiter_data = type(data).__aiter__(data)
+
+            # First fragment.
+            try:
+                # data = anext(aiter_data) without anext
+                data = await type(aiter_data).__anext__(aiter_data)
+            except StopAsyncIteration:
+                return
+            opcode, data = prepare_data(data)
+            await self.write_frame(False, opcode, data)
+
+            # Other fragments.
+            async for data in aiter_data:
+                confirm_opcode, data = prepare_data(data)
+                if confirm_opcode != opcode:
+                    # We're half-way through a fragmented message and we can't
+                    # complete it. This makes the connection unusable.
+                    self.fail_connection(1011)
+                    raise TypeError("data contains inconsistent types")
+                await self.write_frame(False, OP_CONT, data)
+
+            # Final fragment.
+            await self.write_frame(True, OP_CONT, b"")
 
         else:
             raise TypeError("data must be bytes, str, or iterable")


Test patch:
diff --git a/tests/test_protocol.py b/tests/test_protocol.py
index 70c2be0..7a8b0a6 100644
--- a/tests/test_protocol.py
+++ b/tests/test_protocol.py
@@ -27,6 +27,27 @@ if os.environ.get("PYTHONASYNCIODEBUG"):  # pragma: no cover
 MS = max(MS, 2.5 * time.get_clock_info("monotonic").resolution)
 
 
+class async_iterable:
+
+    # In Python ≥ 3.6, this can be simplified to:
+
+    # async def async_iterable(iterable):
+    #     for item in iterable:
+    #         yield item
+
+    def __init__(self, iterable):
+        self.iterator = iter(iterable)
+
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        try:
+            return next(self.iterator)
+        except StopIteration:
+            raise StopAsyncIteration
+
+
 class TransportMock(unittest.mock.Mock):
     """
     Transport mock to control the protocol's inputs and outputs in tests.
@@ -599,6 +620,65 @@ class CommonTests:
             (True, OP_CLOSE, serialize_close(1011, "")),
         )
 
+    def test_send_async_iterable_text(self):
+        self.loop.run_until_complete(self.protocol.send(async_iterable(["ca", "fé"])))
+        self.assertFramesSent(
+            (False, OP_TEXT, "ca".encode("utf-8")),
+            (False, OP_CONT, "fé".encode("utf-8")),
+            (True, OP_CONT, "".encode("utf-8")),
+        )
+
+    def test_send_async_iterable_binary(self):
+        self.loop.run_until_complete(self.protocol.send(async_iterable([b"te", b"a"])))
+        self.assertFramesSent(
+            (False, OP_BINARY, b"te"), (False, OP_CONT, b"a"), (True, OP_CONT, b"")
+        )
+
+    def test_send_async_iterable_binary_from_bytearray(self):
+        self.loop.run_until_complete(
+            self.protocol.send(async_iterable([bytearray(b"te"), bytearray(b"a")]))
+        )
+        self.assertFramesSent(
+            (False, OP_BINARY, b"te"), (False, OP_CONT, b"a"), (True, OP_CONT, b"")
+        )
+
+    def test_send_async_iterable_binary_from_memoryview(self):
+        self.loop.run_until_complete(
+            self.protocol.send(async_iterable([memoryview(b"te"), memoryview(b"a")]))
+        )
+        self.assertFramesSent(
+            (False, OP_BINARY, b"te"), (False, OP_CONT, b"a"), (True, OP_CONT, b"")
+        )
+
+    def test_send_async_iterable_binary_from_non_contiguous_memoryview(self):
+        self.loop.run_until_complete(
+            self.protocol.send(
+                async_iterable([memoryview(b"ttee")[::2], memoryview(b"aa")[::2]])
+            )
+        )
+        self.assertFramesSent(
+            (False, OP_BINARY, b"te"), (False, OP_CONT, b"a"), (True, OP_CONT, b"")
+        )
+
+    def test_send_empty_async_iterable(self):
+        self.loop.run_until_complete(self.protocol.send(async_iterable([])))
+        self.assertNoFrameSent()
+
+    def test_send_async_iterable_type_error(self):
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(self.protocol.send(async_iterable([42])))
+        self.assertNoFrameSent()
+
+    def test_send_async_iterable_mixed_type_error(self):
+        with self.assertRaises(TypeError):
+            self.loop.run_until_complete(
+                self.protocol.send(async_iterable(["café", b"tea"]))
+            )
+        self.assertFramesSent(
+            (False, OP_TEXT, "café".encode("utf-8")),
+            (True, OP_CLOSE, serialize_close(1011, "")),
+        )
+
     def test_send_on_closing_connection_local(self):
         close_task = self.half_close_connection_local()