matthew/enhanced-discord.py
1
0
Fork 0
Code Issues 4 Pull Requests 21 Projects Releases Wiki Activity

Add option to disable auto member chunking.

Browse Source
This commit is contained in:
Rapptz
2017-01-23 07:07:42 -05:00
parent d206ee792b
commit e1aaf74fa7
4 changed files with 156 additions and 97 deletions
Download Patch File Download Diff File Expand all files Collapse all files

92
discord/client.py
View File

@@ -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|

58
discord/shard.py
View File

@@ -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:

89
discord/state.py
View File

@@ -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)
if self._fetch_offline:
guilds = sorted(self._ready_state.guilds, key=lambda g: g.shard_id)
# split chunks by shard ID
chunks = []
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 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)