From 3831485ff3ca7c19f42e5efd1708537732d388c3 Mon Sep 17 00:00:00 2001
From: Pedro Ferreira <pedro@dete.st>
Date: Tue, 1 Aug 2017 22:06:29 +0200
Subject: [PATCH] Add Python 3.5 compatibility

Close #4
---
 secret_handshake/boxstream.py      | 5 ++++-
 secret_handshake/network.py        | 5 ++++-
 secret_handshake/test_boxstream.py | 3 ++-
 secret_handshake/util.py           | 8 ++++++++
 setup.py                           | 1 +
 5 files changed, 19 insertions(+), 3 deletions(-)

diff --git a/secret_handshake/boxstream.py b/secret_handshake/boxstream.py
index 9df3891..af948cf 100644
--- a/secret_handshake/boxstream.py
+++ b/secret_handshake/boxstream.py
@@ -1,6 +1,8 @@
 import struct
 from asyncio import IncompleteReadError
 
+from async_generator import async_generator, yield_
+
 from nacl.secret import SecretBox
 
 from .util import split_chunks, inc_nonce
@@ -58,12 +60,13 @@ class UnboxStream(object):
         self.nonce = inc_nonce(inc_nonce(self.nonce))
         return body
 
+    @async_generator
     async def __aiter__(self):
         while True:
             data = await self.read()
             if data is None:
                 return
-            yield data
+            await yield_(data)
 
 
 class BoxStream(object):
diff --git a/secret_handshake/network.py b/secret_handshake/network.py
index c5fa649..6875bf0 100644
--- a/secret_handshake/network.py
+++ b/secret_handshake/network.py
@@ -21,6 +21,8 @@
 
 from asyncio import open_connection, start_server, ensure_future
 
+from async_generator import async_generator, yield_
+
 from .boxstream import get_stream_pair
 from .crypto import SHSClientCrypto, SHSServerCrypto
 
@@ -43,9 +45,10 @@ class SHSSocket(object):
     def disconnect(self):
         self.writer.close()
 
+    @async_generator
     async def __aiter__(self):
         async for msg in self.read_stream:
-            yield msg
+            await yield_(msg)
 
     def on_connect(self, cb):
         self._on_connect = cb
diff --git a/secret_handshake/test_boxstream.py b/secret_handshake/test_boxstream.py
index fe2db6b..1e7b477 100644
--- a/secret_handshake/test_boxstream.py
+++ b/secret_handshake/test_boxstream.py
@@ -24,6 +24,7 @@ from io import BytesIO
 
 from .test_crypto import (CLIENT_ENCRYPT_KEY, CLIENT_ENCRYPT_NONCE)
 from secret_handshake.boxstream import BoxStream, UnboxStream, HEADER_LENGTH
+from secret_handshake.util import async_comprehend
 
 MESSAGE_1 = (b'\xcev\xedE\x06l\x02\x13\xc8\x17V\xfa\x8bZ?\x88B%O\xb0L\x9f\x8e\x8c0y\x1dv\xc0\xc9\xf6\x9d\xc2\xdf\xdb'
              b'\xee\x9d')
@@ -71,7 +72,7 @@ async def test_unboxstream():
 
     unbox_stream = UnboxStream(buffer, CLIENT_ENCRYPT_KEY, CLIENT_ENCRYPT_NONCE)
     assert not unbox_stream.closed
-    assert [msg async for msg in unbox_stream] == [b'foo', b'foo', b'bar']
+    assert (await async_comprehend(unbox_stream)) == [b'foo', b'foo', b'bar']
     assert unbox_stream.closed
 
 
diff --git a/secret_handshake/util.py b/secret_handshake/util.py
index a98a16d..8972632 100644
--- a/secret_handshake/util.py
+++ b/secret_handshake/util.py
@@ -4,6 +4,14 @@ NONCE_SIZE = 24
 MAX_NONCE = (8 * NONCE_SIZE)
 
 
+async def async_comprehend(generator):
+    """Emulate ``[elem async for elem in generator]``."""
+    results = []
+    async for msg in generator:
+        results.append(msg)
+    return results
+
+
 def inc_nonce(nonce):
     num = bytes_to_long(nonce) + 1
     if num > 2 ** MAX_NONCE:
diff --git a/setup.py b/setup.py
index 04187a4..3850444 100644
--- a/setup.py
+++ b/setup.py
@@ -45,6 +45,7 @@ extras_require = {
 extras_require['all'] = sum((lst for lst in extras_require.values()), [])
 
 install_requires = [
+    'async-generator==1.8',
     'pynacl==1.1.2'
 ]