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

Fixes #494

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.')
 

discord/utils.py

@@ -276,3 +276,10 @@ def async_all(gen):
         if not elem:
             return False
     return True
+
+@asyncio.coroutine
+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()