From 7043ee070207b27b1e36a4ad5c210c689c82db9a Mon Sep 17 00:00:00 2001
From: Paul Sokolovsky <pfalcon@users.sourceforge.net>
Date: Sat, 28 Jan 2017 01:04:21 +0300
Subject: [PATCH] uasyncio: Implement StreamReader.readexactly().

With a unit test.
---
 uasyncio/test_readexactly.py  | 31 +++++++++++++++++++++++++++++++
 uasyncio/uasyncio/__init__.py | 13 +++++++++++++
 2 files changed, 44 insertions(+)
 create mode 100644 uasyncio/test_readexactly.py

diff --git a/uasyncio/test_readexactly.py b/uasyncio/test_readexactly.py
new file mode 100644
index 00000000..59fc53f2
--- /dev/null
+++ b/uasyncio/test_readexactly.py
@@ -0,0 +1,31 @@
+from uasyncio import StreamReader
+
+class MockSock:
+
+    def __init__(self, data_list):
+        self.data = data_list
+
+    def read(self, sz):
+        try:
+            return self.data.pop(0)
+        except IndexError:
+            return b""
+
+
+mock = MockSock([
+    b"123",
+    b"234", b"5",
+    b"a", b"b", b"c", b"d", b"e",
+])
+
+
+def func():
+    sr = StreamReader(mock)
+    assert await sr.readexactly(3) == b"123"
+    assert await sr.readexactly(4) == b"2345"
+    assert await sr.readexactly(5) == b"abcde"
+    # This isn't how it should be, but the current behavior
+    assert await sr.readexactly(10) == b""
+
+for i in func():
+    pass
diff --git a/uasyncio/uasyncio/__init__.py b/uasyncio/uasyncio/__init__.py
index 600dfa0f..14ca59f0 100644
--- a/uasyncio/uasyncio/__init__.py
+++ b/uasyncio/uasyncio/__init__.py
@@ -89,6 +89,19 @@ class StreamReader:
             yield IOReadDone(self.s)
         return res
 
+    def readexactly(self, n):
+        buf = b""
+        while n:
+            yield IORead(self.s)
+            res = self.s.read(n)
+            assert res is not None
+            if not res:
+                yield IOReadDone(self.s)
+                break
+            buf += res
+            n -= len(res)
+        return buf
+
     def readline(self):
         if __debug__:
             log.debug("StreamReader.readline()")