Merge f28f6c9 into e70324e

pferreir · Aug 1, 2017 · c7295f4 · c7295f4
2 parents e70324e + f28f6c9
commit c7295f4
Show file tree

Hide file tree

Showing 2 changed files with 17 additions and 9 deletions.
diff --git a/secret_handshake/boxstream.py b/secret_handshake/boxstream.py
@@ -58,12 +58,14 @@ async def read(self):
         self.nonce = inc_nonce(inc_nonce(self.nonce))
         return body
 
-    async def __aiter__(self):
-        while True:
-            data = await self.read()
-            if data is None:
-                return
-            yield data
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        data = await self.read()
+        if data is None:
+            raise StopAsyncIteration
+        return data
 
 
 class BoxStream(object):

diff --git a/secret_handshake/network.py b/secret_handshake/network.py
@@ -43,9 +43,15 @@ async def read(self):
     def disconnect(self):
         self.writer.close()
 
-    async def __aiter__(self):
-        async for msg in self.read_stream:
-            yield msg
+    def __aiter__(self):
+        return self
+
+    async def __anext__(self):
+        msg = self.read_stream
+        if msg:
+            return msg
+        else:
+            raise StopAsyncIteration
 
     def on_connect(self, cb):
         self._on_connect = cb