Browse Source

Use a proper type for the event queue

pull/5164/head
Rapptz 5 years ago
parent
commit
988505a97f
  1. 43
      discord/shard.py

43
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|

Loading…
Cancel
Save