diff --git a/discord/client.py b/discord/client.py index 25408b7f7..c9ada5df1 100644 --- a/discord/client.py +++ b/discord/client.py @@ -96,6 +96,11 @@ class Client: Integer starting at 0 and less than shard_count. shard_count : Optional[int] The total number of shards. + fetch_offline_members: bool + Indicates if :func:`on_ready` should be delayed to fetch all offline + members from the guilds the bot belongs to. If this is ``False``\, then + no offline members are received and :meth:`request_offline_members` + must be used to fetch the offline members of the guild. Attributes ----------- @@ -120,7 +125,6 @@ class Client: The websocket gateway the client is currently connected to. Could be None. loop The `event loop`_ that the client uses for HTTP requests and websocket operations. - """ def __init__(self, *, loop=None, **options): self.ws = None @@ -133,7 +137,7 @@ class Client: connector = options.pop('connector', None) self.http = HTTPClient(connector, loop=self.loop) - self.connection = ConnectionState(dispatch=self.dispatch, chunker=self.request_offline_members, + self.connection = ConnectionState(dispatch=self.dispatch, chunker=self._chunker, syncer=self._syncer, http=self.http, loop=self.loop, **options) self.connection.shard_count = self.shard_count @@ -151,6 +155,24 @@ class Client: def _syncer(self, guilds): yield from self.ws.request_sync(guilds) + @asyncio.coroutine + def _chunker(self, guild): + if hasattr(guild, 'id'): + guild_id = guild.id + else: + guild_id = [s.id for s in guild] + + payload = { + 'op': 8, + 'd': { + 'guild_id': guild_id, + 'query': '', + 'limit': 0 + } + } + + yield from self.ws.send_as_json(payload) + def handle_reaction_add(self, reaction, user): removed = [] for i, (condition, future, event_type) in enumerate(self._listeners): @@ -261,6 +283,35 @@ class Client: print('Ignoring exception in {}'.format(event_method), file=sys.stderr) traceback.print_exc() + @asyncio.coroutine + def request_offline_members(self, *guilds): + """|coro| + + Requests previously offline members from the guild to be filled up + into the :attr:`Guild.members` cache. This function is usually not + called. It should only be used if you have the ``fetch_offline_members`` + parameter set to ``False``. + + When the client logs on and connects to the websocket, Discord does + not provide the library with offline members if the number of members + in the guild is larger than 250. You can check if a guild is large + if :attr:`Guild.large` is ``True``. + + Parameters + ----------- + \*guilds + An argument list of guilds to request offline members for. + + Raises + ------- + InvalidArgument + If any guild is unavailable or not large in the collection. + """ + if any(not g.large or g.unavailable for g in guilds): + raise InvalidArgument('An unavailable or non-large guild was passed.') + + yield from self.connection.request_offline_members(guilds) + # login state management @asyncio.coroutine @@ -777,43 +828,6 @@ class Client: return self.event(coro) - @asyncio.coroutine - def request_offline_members(self, guild): - """|coro| - - Requests previously offline members from the guild to be filled up - into the :attr:`Guild.members` cache. This function is usually not - called. - - When the client logs on and connects to the websocket, Discord does - not provide the library with offline members if the number of members - in the guild is larger than 250. You can check if a guild is large - if :attr:`Guild.large` is ``True``. - - Parameters - ----------- - guild : :class:`Guild` or iterable - The guild to request offline members for. If this parameter is a - iterable then it is interpreted as an iterator of guilds to - request offline members for. - """ - - if hasattr(guild, 'id'): - guild_id = guild.id - else: - guild_id = [s.id for s in guild] - - payload = { - 'op': 8, - 'd': { - 'guild_id': guild_id, - 'query': '', - 'limit': 0 - } - } - - yield from self.ws.send_as_json(payload) - @asyncio.coroutine def change_presence(self, *, game=None, status=None, afk=False): """|coro| diff --git a/discord/shard.py b/discord/shard.py index 904fd32c4..4ae9768e6 100644 --- a/discord/shard.py +++ b/discord/shard.py @@ -27,13 +27,14 @@ DEALINGS IN THE SOFTWARE. from .state import AutoShardedConnectionState from .client import Client from .gateway import * -from .errors import ConnectionClosed, ClientException +from .errors import ConnectionClosed, ClientException, InvalidArgument from . import compat from .enums import Status import asyncio import logging import websockets +import itertools log = logging.getLogger(__name__) @@ -108,7 +109,7 @@ class AutoShardedClient(Client): elif not isinstance(self.shard_ids, (list, tuple)): raise ClientException('shard_ids parameter must be a list or a tuple.') - self.connection = AutoShardedConnectionState(dispatch=self.dispatch, chunker=self.request_offline_members, + self.connection = AutoShardedConnectionState(dispatch=self.dispatch, chunker=self._chunker, syncer=self._syncer, http=self.http, loop=self.loop, **kwargs) # instead of a single websocket, we have multiple @@ -118,26 +119,7 @@ class AutoShardedClient(Client): self._still_sharding = True @asyncio.coroutine - def request_offline_members(self, guild, *, shard_id=None): - """|coro| - - Requests previously offline members from the guild to be filled up - into the :attr:`Guild.members` cache. This function is usually not - called. - - When the client logs on and connects to the websocket, Discord does - not provide the library with offline members if the number of members - in the guild is larger than 250. You can check if a guild is large - if :attr:`Guild.large` is ``True``. - - Parameters - ----------- - guild: :class:`Guild` or list - The guild to request offline members for. If this parameter is a - list then it is interpreted as a list of guilds to request offline - members for. - """ - + def _chunker(self, guild, *, shard_id=None): try: guild_id = guild.id shard_id = shard_id or guild.shard_id @@ -156,6 +138,38 @@ class AutoShardedClient(Client): ws = self.shards[shard_id].ws yield from ws.send_as_json(payload) + @asyncio.coroutine + def request_offline_members(self, *guilds): + """|coro| + + Requests previously offline members from the guild to be filled up + into the :attr:`Guild.members` cache. This function is usually not + called. It should only be used if you have the ``fetch_offline_members`` + parameter set to ``False``. + + When the client logs on and connects to the websocket, Discord does + not provide the library with offline members if the number of members + in the guild is larger than 250. You can check if a guild is large + if :attr:`Guild.large` is ``True``. + + Parameters + ----------- + \*guilds + An argument list of guilds to request offline members for. + + Raises + ------- + InvalidArgument + If any guild is unavailable or not large in the collection. + """ + if any(not g.large or g.unavailable for g in guilds): + raise InvalidArgument('An unavailable or non-large guild was passed.') + + _guilds = sorted(guilds, key=lambda g: g.shard_id) + for shard_id, sub_guilds in itertools.groupby(_guilds, key=lambda g: g.shard_id): + sub_guilds = list(sub_guilds) + yield from self.connection.request_offline_members(sub_guilds, shard_id=shard_id) + @asyncio.coroutine def pending_reads(self, shard): try: diff --git a/discord/state.py b/discord/state.py index eeb6e894b..6b6a8fc6e 100644 --- a/discord/state.py +++ b/discord/state.py @@ -63,6 +63,7 @@ class ConnectionState: self.syncer = syncer self.is_bot = None self.shard_count = None + self._fetch_offline = options.get('fetch_offline_members', True) self._listeners = [] self.clear() @@ -197,16 +198,7 @@ class ConnectionState: yield self.receive_chunk(guild.id) @asyncio.coroutine - def _delay_ready(self): - 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, loop=self.loop) - - guilds = self._ready_state.guilds - + def request_offline_members(self, guilds): # get all the chunks chunks = [] for guild in guilds: @@ -224,6 +216,22 @@ class ConnectionState: except asyncio.TimeoutError: log.info('Somehow timed out waiting for chunks.') + @asyncio.coroutine + def _delay_ready(self): + launch = self._ready_state.launch + + # only real bots wait for GUILD_CREATE streaming + if self.is_bot: + 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, loop=self.loop) + + guilds = self._ready_state.guilds + if self._fetch_offline: + yield from self.request_offline_members(guilds) + # remove the state try: del self._ready_state @@ -260,6 +268,7 @@ class ConnectionState: factory, _ = _channel_factory(pm['type']) self._add_private_channel(factory(me=self.user, data=pm, state=self)) + self.dispatch('connect') compat.create_task(self._delay_ready(), loop=self.loop) def parse_resumed(self, data): @@ -477,8 +486,8 @@ class ConnectionState: @asyncio.coroutine def _chunk_and_dispatch(self, guild, unavailable): - yield from self.chunker(guild) chunks = list(self.chunks_needed(guild)) + yield from self.chunker(guild) if chunks: try: yield from asyncio.wait(chunks, timeout=len(chunks), loop=self.loop) @@ -518,9 +527,10 @@ class ConnectionState: return # since we're not waiting for 'useful' READY we'll just - # do the chunk request here - compat.create_task(self._chunk_and_dispatch(guild, unavailable), loop=self.loop) - return + # do the chunk request here if wanted + if self._fetch_offline: + compat.create_task(self._chunk_and_dispatch(guild, unavailable), loop=self.loop) + return # Dispatch available if newly available if unavailable == False: @@ -740,6 +750,25 @@ class AutoShardedConnectionState(ConnectionState): self._ready_state = ReadyState(launch=asyncio.Event(), guilds=[]) self._ready_task = None + @asyncio.coroutine + def request_offline_members(self, guilds, *, shard_id): + # get all the chunks + chunks = [] + for guild in guilds: + chunks.extend(self.chunks_needed(guild)) + + # we only want to request ~75 guilds per chunk request. + splits = [guilds[i:i + 75] for i in range(0, len(guilds), 75)] + for split in splits: + yield from self.chunker(split, shard_id=shard_id) + + # wait for the chunks + if chunks: + try: + yield from asyncio.wait(chunks, timeout=len(chunks) * 30.0, loop=self.loop) + except asyncio.TimeoutError: + log.info('Somehow timed out waiting for chunks.') + @asyncio.coroutine def _delay_ready(self): launch = self._ready_state.launch @@ -749,30 +778,14 @@ class AutoShardedConnectionState(ConnectionState): launch.set() yield from asyncio.sleep(2.0 * self.shard_count, loop=self.loop) - guilds = sorted(self._ready_state.guilds, key=lambda g: g.shard_id) - - # we only want to request ~75 guilds per chunk request. - # we also want to split the chunks per shard_id - for shard_id, sub_guilds in itertools.groupby(guilds, key=lambda g: g.shard_id): - sub_guilds = list(sub_guilds) - # split chunks by shard ID - chunks = [] - for guild in sub_guilds: - chunks.extend(self.chunks_needed(guild)) + if self._fetch_offline: + guilds = sorted(self._ready_state.guilds, key=lambda g: g.shard_id) - splits = [sub_guilds[i:i + 75] for i in range(0, len(sub_guilds), 75)] - for split in splits: - yield from self.chunker(split, shard_id=shard_id) - - # wait for the chunks - if chunks: - try: - yield from asyncio.wait(chunks, timeout=len(chunks) * 30.0, loop=self.loop) - except asyncio.TimeoutError: - log.info('Somehow timed out waiting for chunks for %s shard_id' % shard_id) - - self.dispatch('shard_ready', shard_id) + for shard_id, sub_guilds in itertools.groupby(guilds, key=lambda g: g.shard_id): + sub_guilds = list(sub_guilds) + yield from self.request_offline_members(sub_guilds, shard_id=shard_id) + self.dispatch('shard_ready', shard_id) # remove the state try: @@ -782,6 +795,9 @@ class AutoShardedConnectionState(ConnectionState): # regular users cannot shard so we won't worry about it here. + # clear the current task + self._ready_task = None + # dispatch the event self.dispatch('ready') @@ -801,5 +817,6 @@ class AutoShardedConnectionState(ConnectionState): factory, _ = _channel_factory(pm['type']) self._add_private_channel(factory(me=self.user, data=pm, state=self)) + self.dispatch('connect') if self._ready_task is None: self._ready_task = compat.create_task(self._delay_ready(), loop=self.loop) diff --git a/docs/api.rst b/docs/api.rst index f8e0e8b5f..b31fa280a 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -102,6 +102,13 @@ to handle it, which defaults to print a traceback and ignore the exception. .. versionadded:: 0.7.0 Subclassing to listen to events. +.. function:: on_connect() + + Called when the client has successfully connected to Discord. This is not + the same as the client being fully prepared, see :func:`on_ready` for that. + + The warnings on :func:`on_ready` also apply. + .. function:: on_ready() Called when the client is done preparing the data received from Discord. Usually after login is successful @@ -114,6 +121,13 @@ to handle it, which defaults to print a traceback and ignore the exception. once. This library implements reconnection logic and thus will end up calling this event whenever a RESUME request fails. +.. function:: on_shard_ready(shard_id) + + Similar to :func:`on_ready` except used by :class:`AutoShardedClient` + to denote when a particular shard ID has become ready. + + :param shard_id: The shard ID that is ready. + .. function:: on_resumed() Called when the client has resumed a session.