Wrap asyncio.wait into a saner alternative that raises TimeoutError.

Fixes #494

2 files changed, 10 insertions, 3 deletions

diff --git a/discord/state.py b/discord/state.py
index 8e1ba0eb..d0e0a71f 100644
--- a/discord/state.py
+++ b/discord/state.py
@@ -216,7 +216,7 @@ class ConnectionState:
         # wait for the chunks
         if chunks:
             try:
-                yield from asyncio.wait(chunks, timeout=len(chunks) * 30.0, loop=self.loop)
+                yield from utils.sane_wait_for(chunks, timeout=len(chunks) * 30.0, loop=self.loop)
             except asyncio.TimeoutError:
                 log.info('Somehow timed out waiting for chunks.')
 
@@ -494,7 +494,7 @@ class ConnectionState:
         yield from self.chunker(guild)
         if chunks:
             try:
-                yield from asyncio.wait(chunks, timeout=len(chunks), loop=self.loop)
+                yield from utils.sane_wait_for(chunks, timeout=len(chunks), loop=self.loop)
             except asyncio.TimeoutError:
                 log.info('Somehow timed out waiting for chunks.')
 
@@ -773,7 +773,7 @@ class AutoShardedConnectionState(ConnectionState):
         # wait for the chunks
         if chunks:
             try:
-                yield from asyncio.wait(chunks, timeout=len(chunks) * 30.0, loop=self.loop)
+                yield from utils.sane_wait_for(chunks, timeout=len(chunks) * 30.0, loop=self.loop)
             except asyncio.TimeoutError:
                 log.info('Somehow timed out waiting for chunks.')
 
diff --git a/discord/utils.py b/discord/utils.py
index 5177b1c8..2c5ddc90 100644
--- a/discord/utils.py
+++ b/discord/utils.py
@@ -276,3 +276,10 @@ def async_all(gen):
         if not elem:
             return False
     return True
+
[email protected]
+def sane_wait_for(futures, *, timeout, loop):
+    done, pending = yield from asyncio.wait(futures, timeout=timeout, loop=loop)
+
+    if len(pending) != 0:
+        raise asyncio.TimeoutError()