Use a proper type for the event queue

1 files changed, 28 insertions, 15 deletions

diff --git a/discord/shard.py b/discord/shard.py
index 1e34a56c..0f63b65c 100644
--- a/discord/shard.py
+++ b/discord/shard.py
@@ -42,6 +42,27 @@ class EventType:
     resume = 1
     identify = 2
 
+class EventItem:
+    __slots__ = ('type', 'shard', 'error')
+
+    def __init__(self, etype, shard, error):
+        self.type = etype
+        self.shard = shard
+        self.error = error
+
+    def __lt__(self, other):
+        if not isinstance(other, EventItem):
+            return NotImplemented
+        return self.type < other.type
+
+    def __eq__(self, other):
+        if not isinstance(other, EventItem):
+            return NotImplemented
+        return self.type == other.type
+
+    def __hash__(self):
+        return hash(self.type)
+
 class Shard:
     def __init__(self, ws, client):
         self.ws = ws
@@ -64,10 +85,10 @@ class Shard:
                 await self.ws.poll_event()
             except ReconnectWebSocket as e:
                 etype = EventType.resume if e.resume else EventType.identify
-                self._queue.put_nowait((etype, self, e))
+                self._queue.put_nowait(EventItem(etype, self, e))
                 break
             except ConnectionClosed as e:
-                self._queue.put_nowait((EventType.close, self, e))
+                self._queue.put_nowait(EventItem(EventType.close, self, e))
                 break
 
     async def reconnect(self, exc):
@@ -220,23 +241,15 @@ class AutoShardedClient(Client):
             if shard_id != last_shard_id:
                 await asyncio.sleep(5.0)
 
-        # shards_to_wait_for = []
-        # for shard in self.shards.values():
-        #     shard.complete_pending_reads()
-        #     shards_to_wait_for.append(shard.wait())
-
-        # # wait for all pending tasks to finish
-        # await utils.sane_wait_for(shards_to_wait_for, timeout=300.0)
-
     async def _connect(self):
         await self.launch_shards()
 
         while True:
-            etype, shard, exc = await self._queue.get()
-            if etype == EventType.close:
-                raise exc
-            elif etype in (EventType.identify, EventType.resume):
-                await shard.reconnect(exc)
+            item = await self._queue.get()
+            if item.type == EventType.close:
+                raise item.error
+            elif item.type in (EventType.identify, EventType.resume):
+                await item.shard.reconnect(item.error)
 
     async def close(self):
         """|coro|