From 425bd2c0911cd90aa30ab18e79b60078ee1cb0f0 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Sun, 6 Mar 2016 05:24:18 -0500 Subject: Move chunking logic back into ConnectionState. This allows for a nicer design when dealing with parsers that could end up being coroutines. --- discord/state.py | 24 +++++++++++++++++++++--- 1 file changed, 21 insertions(+), 3 deletions(-) (limited to 'discord/state.py') diff --git a/discord/state.py b/discord/state.py index 6db17218..fec2e111 100644 --- a/discord/state.py +++ b/discord/state.py @@ -36,10 +36,9 @@ from .enums import Status from collections import deque, namedtuple -import copy +import copy, enum, math import datetime import asyncio -import enum import logging class ListenerType(enum.Enum): @@ -49,10 +48,11 @@ Listener = namedtuple('Listener', ('type', 'future', 'predicate')) log = logging.getLogger(__name__) class ConnectionState: - def __init__(self, dispatch, max_messages, *, loop): + def __init__(self, dispatch, chunker, max_messages, *, loop): self.loop = loop self.max_messages = max_messages self.dispatch = dispatch + self.chunker = chunker self._listeners = [] self.clear() @@ -128,6 +128,7 @@ class ConnectionState: self._add_server(server) return server + @asyncio.coroutine def parse_ready(self, data): self.user = User(**data['user']) guilds = data.get('guilds') @@ -139,6 +140,23 @@ class ConnectionState: self._add_private_channel(PrivateChannel(id=pm['id'], user=User(**pm['recipient']))) + # 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_needed = math.ceil(server._member_count / 1000) + for chunk in range(chunks_needed): + chunks.append(self.receive_chunk(server.id)) + + if chunks: + yield from asyncio.wait(chunks) + + self.dispatch('ready') + def parse_message_create(self, data): channel = self.get_channel(data.get('channel_id')) message = Message(channel=channel, **data) -- cgit v1.2.3