From cb8cb557f539baf41640cad3037be02ff9b79818 Mon Sep 17 00:00:00 2001 From: Rapptz Date: Thu, 10 Sep 2020 05:56:48 -0400 Subject: [PATCH] Add Guild.chunk and deprecated Client.request_offline_members --- discord/client.py | 3 ++- discord/guild.py | 26 ++++++++++++++++----- discord/shard.py | 3 ++- discord/state.py | 57 ++++++++++++++++++++++++++--------------------- 4 files changed, 57 insertions(+), 32 deletions(-) diff --git a/discord/client.py b/discord/client.py index dc8b360d7..c39a26da4 100644 --- a/discord/client.py +++ b/discord/client.py @@ -382,6 +382,7 @@ class Client: print('Ignoring exception in {}'.format(event_method), file=sys.stderr) traceback.print_exc() + @utils.deprecated('Guild.chunk') async def request_offline_members(self, *guilds): r"""|coro| @@ -397,7 +398,7 @@ class Client: .. warning:: - This method is deprecated. + This method is deprecated. Use :meth:`Guild.chunk` instead. Parameters ----------- diff --git a/discord/guild.py b/discord/guild.py index c15b78cf4..fcc4e4b30 100644 --- a/discord/guild.py +++ b/discord/guild.py @@ -2037,6 +2037,23 @@ class Guild(Hashable): return Widget(state=self._state, data=data) + async def chunk(self, *, cache=True): + """|coro| + + Requests all members that belong to this guild. In order to use this, + :meth:`Intents.members` must be enabled. + + This is a websocket operation and can be slow. + + .. versionadded:: 1.5 + + Parameters + ----------- + cache: :class:`bool` + Whether to cache the members as well. + """ + return await self._state.chunk_guild(self, cache=cache) + async def query_members(self, query=None, *, limit=5, user_ids=None, cache=True): """|coro| @@ -2049,16 +2066,15 @@ class Guild(Hashable): Parameters ----------- - query: :class:`str` - The string that the username's start with. An empty string - requests all members. + query: Optional[:class:`str`] + The string that the username's start with. limit: :class:`int` The maximum number of members to send back. This must be - a number between 1 and 100. + a number between 5 and 100. cache: :class:`bool` Whether to cache the members internally. This makes operations such as :meth:`get_member` work for those that matched. - user_ids: List[:class:`int`] + user_ids: Optional[List[:class:`int`]] List of user IDs to search for. If the user ID is not in the guild then it won't be returned. .. versionadded:: 1.4 diff --git a/discord/shard.py b/discord/shard.py index 1b635c1c0..00a7a1178 100644 --- a/discord/shard.py +++ b/discord/shard.py @@ -334,6 +334,7 @@ class AutoShardedClient(Client): """Mapping[int, :class:`ShardInfo`]: Returns a mapping of shard IDs to their respective info object.""" return { shard_id: ShardInfo(parent, self.shard_count) for shard_id, parent in self.__shards.items() } + @utils.deprecated('Guild.chunk') async def request_offline_members(self, *guilds): r"""|coro| @@ -349,7 +350,7 @@ class AutoShardedClient(Client): .. warning:: - This method is deprecated. + This method is deprecated. Use :meth:`Guild.chunk` instead. Parameters ----------- diff --git a/discord/state.py b/discord/state.py index 088633363..ad491197d 100644 --- a/discord/state.py +++ b/discord/state.py @@ -57,12 +57,28 @@ from .object import Object from .invite import Invite class ChunkRequest: - __slots__ = ('guild_id', 'nonce', 'future') - - def __init__(self, guild_id, future): + def __init__(self, guild_id, future, resolver, *, cache=True): self.guild_id = guild_id + self.resolver = resolver + self.cache = cache self.nonce = os.urandom(16).hex() self.future = future + self.buffer = [] # List[Member] + + def add_members(self, members): + self.buffer.extend(members) + if self.cache: + guild = self.resolver(self.guild_id) + if guild is None: + return + + for member in members: + existing = guild.get_member(member.id) + if existing is None or existing.joined_at is None: + guild._add_member(member) + + def done(self): + self.future.set_result(self.buffer) log = logging.getLogger(__name__) @@ -156,7 +172,7 @@ class ConnectionState: # to reconnect loops which cause mass allocations and deallocations. gc.collect() - def process_chunk_requests(self, guild_id, nonce, members): + def process_chunk_requests(self, guild_id, nonce, members, complete): removed = [] for i, request in enumerate(self._chunk_requests): future = request.future @@ -165,8 +181,10 @@ class ConnectionState: continue if request.guild_id == guild_id and request.nonce == nonce: - future.set_result(members) - removed.append(i) + request.add_members(members) + if complete: + request.done() + removed.append(i) for index in reversed(removed): del self._chunk_requests[index] @@ -330,19 +348,13 @@ class ConnectionState: raise RuntimeError('Somehow do not have a websocket for this guild_id') future = self.loop.create_future() - request = ChunkRequest(guild.id, future) + request = ChunkRequest(guild.id, future, self._get_guild, cache=cache) self._chunk_requests.append(request) try: # start the query operation await ws.request_chunks(guild_id, query=query, limit=limit, user_ids=user_ids, nonce=request.nonce) - members = await asyncio.wait_for(future, timeout=30.0) - - if cache: - for member in members: - guild._add_member(member) - - return members + return await asyncio.wait_for(future, timeout=30.0) except asyncio.TimeoutError: log.warning('Timed out waiting for chunks with query %r and limit %d for guild_id %d', query, limit, guild_id) raise @@ -747,9 +759,10 @@ class ConnectionState: return self._add_guild_from_data(data) - async def chunk_guild(self, guild, *, wait=True): + async def chunk_guild(self, guild, *, wait=True, cache=None): + cache = cache or self._cache_members future = self.loop.create_future() - request = ChunkRequest(guild.id, future) + request = ChunkRequest(guild.id, future, self._get_guild, cache=cache) self._chunk_requests.append(request) await self.chunker(guild.id, nonce=request.nonce) if wait: @@ -893,14 +906,8 @@ class ConnectionState: guild = self._get_guild(guild_id) members = [Member(guild=guild, data=member, state=self) for member in data.get('members', [])] log.debug('Processed a chunk for %s members in guild ID %s.', len(members), guild_id) - if self._cache_members: - for member in members: - existing = guild.get_member(member.id) - if existing is None or existing.joined_at is None: - guild._add_member(member) - - if data.get('chunk_index', 0) + 1 == data.get('chunk_count'): - self.process_chunk_requests(guild_id, data.get('nonce'), members) + complete = data.get('chunk_index', 0) + 1 == data.get('chunk_count') + self.process_chunk_requests(guild_id, data.get('nonce'), members, complete) def parse_guild_integrations_update(self, data): guild = self._get_guild(int(data['guild_id'])) @@ -1075,7 +1082,7 @@ class AutoShardedConnectionState(ConnectionState): current_bucket.append(future) else: future = self.loop.create_future() - future.set_result(True) + future.set_result([]) processed.append((guild, future))