Add option to disable auto member chunking.

4 files changed, 156 insertions, 97 deletions

diff --git a/discord/client.py b/discord/client.py
index 25408b7f..c9ada5df 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
@@ -778,43 +829,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 904fd32c..4ae9768e 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
@@ -157,6 +139,38 @@ class AutoShardedClient(Client):
         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:
             while self._still_sharding:
diff --git a/discord/state.py b/discord/state.py
index eeb6e894..6b6a8fc6 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:
@@ -741,6 +751,25 @@ class AutoShardedConnectionState(ConnectionState):
         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
         while not launch.is_set():
@@ -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 f8e0e8b5..b31fa280 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.