diff options
| -rw-r--r-- | discord/state.py | 90 |
1 files changed, 66 insertions, 24 deletions
diff --git a/discord/state.py b/discord/state.py index fa296f82..4baab8b6 100644 --- a/discord/state.py +++ b/discord/state.py @@ -46,6 +46,7 @@ class ListenerType(enum.Enum): Listener = namedtuple('Listener', ('type', 'future', 'predicate')) log = logging.getLogger(__name__) +ReadyState = namedtuple('ReadyState', ('launch', 'chunks')) class ConnectionState: def __init__(self, dispatch, chunker, max_messages, *, loop): @@ -133,33 +134,50 @@ class ConnectionState: yield self.receive_chunk(server.id) @asyncio.coroutine - def _fill_offline(self): - # a chunk has a maximum of 1000 members. - # we need to find out how many futures we're actually waiting for - large_servers = [s for s in self.servers if s.large] - yield from self.chunker(large_servers) - - chunks = [] - for server in large_servers: - chunks.extend(self.chunks_needed(server)) - + def _delay_ready(self, large_servers): + if len(large_servers): + # for regular accounts with < 100 guilds you will + # get a regular READY packet without the unavailable + # streaming so if this is non-empty then it's a regular + # account and it needs chunking. + yield from self.chunker(large_servers) + for server in large_servers: + self._ready_state.chunks.extend(self.chunks_needed(server)) + + launch = self._ready_state.launch + while not launch.is_set(): + # this snippet of code is basically waiting 2 seconds + # until the last GUILD_CREATE was sent + launch.set() + yield from asyncio.sleep(2) + + # get all the chunks + chunks = [f for f in self._ready_state.chunks if not f.done()] if chunks: yield from asyncio.wait(chunks) + # remove the state + del self._ready_state + + # dispatch the event self.dispatch('ready') def parse_ready(self, data): + self._ready_state = ReadyState(launch=asyncio.Event(), chunks=[]) self.user = User(**data['user']) guilds = data.get('guilds') + large_servers = [] for guild in guilds: - self._add_server_from_data(guild) + server = self._add_server_from_data(guild) + if server.large: + large_servers.append(server) for pm in data.get('private_channels'): self._add_private_channel(PrivateChannel(id=pm['id'], user=User(**pm['recipient']))) - utils.create_task(self._fill_offline(), loop=self.loop) + utils.create_task(self._delay_ready(large_servers), loop=self.loop) def parse_message_create(self, data): channel = self.get_channel(data.get('channel_id')) @@ -295,10 +313,8 @@ class ConnectionState: self.dispatch('member_update', old_member, member) - @asyncio.coroutine - def parse_guild_create(self, data): - unavailable = data.get('unavailable') - if unavailable == False: + def _get_create_server(self, data): + if data.get('unavailable') == False: # GUILD_CREATE with unavailable in the response # usually means that the server has become available # and is therefore in the cache @@ -306,27 +322,53 @@ class ConnectionState: if server is not None: server.unavailable = False server._from_data(data) - self.dispatch('server_available', server) - return + return server + + return self._add_server_from_data(data) + @asyncio.coroutine + def parse_guild_create(self, data): + unavailable = data.get('unavailable') if unavailable == True: # joined a server with unavailable == True so.. return - # if we're at this point then it was probably - # unavailable during the READY event and is now - # available, so it isn't in the cache... - - server = self._add_server_from_data(data) + server = self._get_create_server(data) # check if it requires chunking if server.large: yield from self.chunker(server) chunks = list(self.chunks_needed(server)) + + if unavailable == False: + # check if we're waiting for 'useful' READY + # and if we are, we don't want to dispatch any + # event such as server_join or server_available + # because we're still in the 'READY' phase. Or + # so we say. + try: + state = self._ready_state + state.launch.clear() + if chunks: + state.chunks.extend(chunks) + except AttributeError: + # the _ready_state attribute is only there during + # processing of useful READY. + pass + else: + return + + # since we're not waiting for 'useful' READY we'll just + # do the chunk request here if chunks: yield from asyncio.wait(chunks) - self.dispatch('server_join', server) + + # Dispatch available if newly available + if unavailable == False: + self.dispatch('server_available', server) + else: + self.dispatch('server_join', server) def parse_guild_update(self, data): server = self._get_server(data.get('id')) |