Merge 72754ebed1 into d11ca092f7

docs/library/asyncio.rst

@@ -277,11 +277,27 @@ TCP stream connections
 
     This is a coroutine.
 
+.. method:: Stream.awrite(buf)
+
+    Write *buf* to the stream.
+
+    This method is slightly more efficient than the usual
+    `Stream.write`-plus-`Stream.drain` combination. It ignores the output
+    buffer.
+
+    This is a coroutine, and a MicroPython extension.
+
 .. method:: Stream.write(buf)
 
-    Accumulated *buf* to the output buffer.  The data is only flushed when
-    `Stream.drain` is called.  It is recommended to call `Stream.drain` immediately
-    after calling this function.
+    Add *buf* to the output buffer.  The data is only flushed when
+    `Stream.drain` is called.  It is recommended to call `Stream.drain`
+    immediately after calling this function.
+
+    If the output buffer is empty when *write* is called, part or all of
+    *buf* might be written immediately instead of getting buffered.
+
+    If compatibility to CPython is not required, consider using `Stream.awrite`
+    instead.
 
 .. method:: Stream.drain()
 

extmod/asyncio/stream.py

@@ -73,19 +73,22 @@ class Stream:
                 buf = buf[ret:]
         self.out_buf += buf
 
-    # async
-    def drain(self):
-        if not self.out_buf:
-            # Drain must always yield, so a tight loop of write+drain can't block the scheduler.
-            return (yield from core.sleep_ms(0))
-        mv = memoryview(self.out_buf)
+    async def awrite(self, buf):
+        mv = memoryview(buf)
         off = 0
         while off < len(mv):
             yield core._io_queue.queue_write(self.s)
             ret = self.s.write(mv[off:])
             if ret is not None:
                 off += ret
+
+    async def drain(self):
+        buf = self.out_buf
+        if not buf:
+            # Drain must always yield, so a tight loop of write+drain can't block the scheduler.
+            return await core.sleep_ms(0)
         self.out_buf = b""
+        await self.awrite(buf)
 
 
 # Stream can be used for both reading and writing to save code size

tests/multi_net/asyncio_tcp_readall.py

@@ -10,17 +10,13 @@ PORT = 8000
 
 
 async def handle_connection(reader, writer):
-    writer.write(b"a")
-    await writer.drain()
+    await writer.awrite(b"a")
 
     # Split the first 2 bytes up so the client must wait for the second one
     await asyncio.sleep(1)
 
-    writer.write(b"b")
-    await writer.drain()
-
-    writer.write(b"c")
-    await writer.drain()
+    await writer.awrite(b"b")
+    await writer.awrite(b"c")
 
     print("close")
     writer.close()