Browse Source

Add option to disable auto member chunking.

pull/468/head
Rapptz 8 years ago
parent
commit
e1aaf74fa7
  1. 92
      discord/client.py
  2. 58
      discord/shard.py
  3. 89
      discord/state.py
  4. 14
      docs/api.rst

92
discord/client.py

@ -96,6 +96,11 @@ class Client:
Integer starting at 0 and less than shard_count. Integer starting at 0 and less than shard_count.
shard_count : Optional[int] shard_count : Optional[int]
The total number of shards. 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 Attributes
----------- -----------
@ -120,7 +125,6 @@ class Client:
The websocket gateway the client is currently connected to. Could be None. The websocket gateway the client is currently connected to. Could be None.
loop loop
The `event loop`_ that the client uses for HTTP requests and websocket operations. The `event loop`_ that the client uses for HTTP requests and websocket operations.
""" """
def __init__(self, *, loop=None, **options): def __init__(self, *, loop=None, **options):
self.ws = None self.ws = None
@ -133,7 +137,7 @@ class Client:
connector = options.pop('connector', None) connector = options.pop('connector', None)
self.http = HTTPClient(connector, loop=self.loop) 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) syncer=self._syncer, http=self.http, loop=self.loop, **options)
self.connection.shard_count = self.shard_count self.connection.shard_count = self.shard_count
@ -151,6 +155,24 @@ class Client:
def _syncer(self, guilds): def _syncer(self, guilds):
yield from self.ws.request_sync(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): def handle_reaction_add(self, reaction, user):
removed = [] removed = []
for i, (condition, future, event_type) in enumerate(self._listeners): 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) print('Ignoring exception in {}'.format(event_method), file=sys.stderr)
traceback.print_exc() 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 # login state management
@asyncio.coroutine @asyncio.coroutine
@ -777,43 +828,6 @@ class Client:
return self.event(coro) 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 @asyncio.coroutine
def change_presence(self, *, game=None, status=None, afk=False): def change_presence(self, *, game=None, status=None, afk=False):
"""|coro| """|coro|

58
discord/shard.py

@ -27,13 +27,14 @@ DEALINGS IN THE SOFTWARE.
from .state import AutoShardedConnectionState from .state import AutoShardedConnectionState
from .client import Client from .client import Client
from .gateway import * from .gateway import *
from .errors import ConnectionClosed, ClientException from .errors import ConnectionClosed, ClientException, InvalidArgument
from . import compat from . import compat
from .enums import Status from .enums import Status
import asyncio import asyncio
import logging import logging
import websockets import websockets
import itertools
log = logging.getLogger(__name__) log = logging.getLogger(__name__)
@ -108,7 +109,7 @@ class AutoShardedClient(Client):
elif not isinstance(self.shard_ids, (list, tuple)): elif not isinstance(self.shard_ids, (list, tuple)):
raise ClientException('shard_ids parameter must be a list or a 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) syncer=self._syncer, http=self.http, loop=self.loop, **kwargs)
# instead of a single websocket, we have multiple # instead of a single websocket, we have multiple
@ -118,26 +119,7 @@ class AutoShardedClient(Client):
self._still_sharding = True self._still_sharding = True
@asyncio.coroutine @asyncio.coroutine
def request_offline_members(self, guild, *, shard_id=None): def _chunker(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.
"""
try: try:
guild_id = guild.id guild_id = guild.id
shard_id = shard_id or guild.shard_id shard_id = shard_id or guild.shard_id
@ -156,6 +138,38 @@ class AutoShardedClient(Client):
ws = self.shards[shard_id].ws ws = self.shards[shard_id].ws
yield from ws.send_as_json(payload) 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 @asyncio.coroutine
def pending_reads(self, shard): def pending_reads(self, shard):
try: try:

89
discord/state.py

@ -63,6 +63,7 @@ class ConnectionState:
self.syncer = syncer self.syncer = syncer
self.is_bot = None self.is_bot = None
self.shard_count = None self.shard_count = None
self._fetch_offline = options.get('fetch_offline_members', True)
self._listeners = [] self._listeners = []
self.clear() self.clear()
@ -197,16 +198,7 @@ class ConnectionState:
yield self.receive_chunk(guild.id) yield self.receive_chunk(guild.id)
@asyncio.coroutine @asyncio.coroutine
def _delay_ready(self): def request_offline_members(self, guilds):
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
# get all the chunks # get all the chunks
chunks = [] chunks = []
for guild in guilds: for guild in guilds:
@ -224,6 +216,22 @@ class ConnectionState:
except asyncio.TimeoutError: except asyncio.TimeoutError:
log.info('Somehow timed out waiting for chunks.') 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 # remove the state
try: try:
del self._ready_state del self._ready_state
@ -260,6 +268,7 @@ class ConnectionState:
factory, _ = _channel_factory(pm['type']) factory, _ = _channel_factory(pm['type'])
self._add_private_channel(factory(me=self.user, data=pm, state=self)) self._add_private_channel(factory(me=self.user, data=pm, state=self))
self.dispatch('connect')
compat.create_task(self._delay_ready(), loop=self.loop) compat.create_task(self._delay_ready(), loop=self.loop)
def parse_resumed(self, data): def parse_resumed(self, data):
@ -477,8 +486,8 @@ class ConnectionState:
@asyncio.coroutine @asyncio.coroutine
def _chunk_and_dispatch(self, guild, unavailable): def _chunk_and_dispatch(self, guild, unavailable):
yield from self.chunker(guild)
chunks = list(self.chunks_needed(guild)) chunks = list(self.chunks_needed(guild))
yield from self.chunker(guild)
if chunks: if chunks:
try: try:
yield from asyncio.wait(chunks, timeout=len(chunks), loop=self.loop) yield from asyncio.wait(chunks, timeout=len(chunks), loop=self.loop)
@ -518,9 +527,10 @@ class ConnectionState:
return return
# since we're not waiting for 'useful' READY we'll just # since we're not waiting for 'useful' READY we'll just
# do the chunk request here # do the chunk request here if wanted
compat.create_task(self._chunk_and_dispatch(guild, unavailable), loop=self.loop) if self._fetch_offline:
return compat.create_task(self._chunk_and_dispatch(guild, unavailable), loop=self.loop)
return
# Dispatch available if newly available # Dispatch available if newly available
if unavailable == False: if unavailable == False:
@ -740,6 +750,25 @@ class AutoShardedConnectionState(ConnectionState):
self._ready_state = ReadyState(launch=asyncio.Event(), guilds=[]) self._ready_state = ReadyState(launch=asyncio.Event(), guilds=[])
self._ready_task = None 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 @asyncio.coroutine
def _delay_ready(self): def _delay_ready(self):
launch = self._ready_state.launch launch = self._ready_state.launch
@ -749,30 +778,14 @@ class AutoShardedConnectionState(ConnectionState):
launch.set() launch.set()
yield from asyncio.sleep(2.0 * self.shard_count, loop=self.loop) 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 if self._fetch_offline:
chunks = [] guilds = sorted(self._ready_state.guilds, key=lambda g: g.shard_id)
for guild in sub_guilds:
chunks.extend(self.chunks_needed(guild))
splits = [sub_guilds[i:i + 75] for i in range(0, len(sub_guilds), 75)] for shard_id, sub_guilds in itertools.groupby(guilds, key=lambda g: g.shard_id):
for split in splits: sub_guilds = list(sub_guilds)
yield from self.chunker(split, shard_id=shard_id) yield from self.request_offline_members(sub_guilds, shard_id=shard_id)
self.dispatch('shard_ready', 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)
# remove the state # remove the state
try: try:
@ -782,6 +795,9 @@ class AutoShardedConnectionState(ConnectionState):
# regular users cannot shard so we won't worry about it here. # regular users cannot shard so we won't worry about it here.
# clear the current task
self._ready_task = None
# dispatch the event # dispatch the event
self.dispatch('ready') self.dispatch('ready')
@ -801,5 +817,6 @@ class AutoShardedConnectionState(ConnectionState):
factory, _ = _channel_factory(pm['type']) factory, _ = _channel_factory(pm['type'])
self._add_private_channel(factory(me=self.user, data=pm, state=self)) self._add_private_channel(factory(me=self.user, data=pm, state=self))
self.dispatch('connect')
if self._ready_task is None: if self._ready_task is None:
self._ready_task = compat.create_task(self._delay_ready(), loop=self.loop) self._ready_task = compat.create_task(self._delay_ready(), loop=self.loop)

14
docs/api.rst

@ -102,6 +102,13 @@ to handle it, which defaults to print a traceback and ignore the exception.
.. versionadded:: 0.7.0 .. versionadded:: 0.7.0
Subclassing to listen to events. 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() .. function:: on_ready()
Called when the client is done preparing the data received from Discord. Usually after login is successful 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 once. This library implements reconnection logic and thus will
end up calling this event whenever a RESUME request fails. 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() .. function:: on_resumed()
Called when the client has resumed a session. Called when the client has resumed a session.

Loading…
Cancel
Save