From 5b2f630848f42940d62ce7af5ee1484a39259ceb Mon Sep 17 00:00:00 2001 From: Rapptz Date: Mon, 15 Jul 2019 07:56:35 -0400 Subject: [PATCH] Add Guild.query_members to fetch members from the gateway. --- discord/gateway.py | 11 ++++++++++ discord/guild.py | 40 +++++++++++++++++++++++++++++++++++++ discord/state.py | 50 +++++++++++++++++++++++++++++++++++++++------- discord/utils.py | 4 +++- 4 files changed, 97 insertions(+), 8 deletions(-) diff --git a/discord/gateway.py b/discord/gateway.py index 7d71425be..6848f120f 100644 --- a/discord/gateway.py +++ b/discord/gateway.py @@ -517,6 +517,17 @@ class DiscordWebSocket(websockets.client.WebSocketClientProtocol): } await self.send_as_json(payload) + async def request_chunks(self, guild_id, query, limit): + payload = { + 'op': self.REQUEST_MEMBERS, + 'd': { + 'guild_id': str(guild_id), + 'query': query, + 'limit': limit + } + } + await self.send_as_json(payload) + async def voice_state(self, guild_id, channel_id, self_mute=False, self_deaf=False): payload = { 'op': self.VOICE_STATE, diff --git a/discord/guild.py b/discord/guild.py index 4c87e9f02..5a95f5dac 100644 --- a/discord/guild.py +++ b/discord/guild.py @@ -1844,3 +1844,43 @@ class Guild(Hashable): data = await self._state.http.get_widget(self.id) return Widget(state=self._state, data=data) + + async def query_members(self, query, *, limit=5, cache=True): + """|coro| + + Request members that belong to this guild whose username starts with + the query given. + + This is a websocket operation and can be slow. + + .. warning:: + + Most bots do not need to use this. It's mainly a helper + for bots who have disabled ``guild_subscriptions``. + + .. versionadded:: 1.3 + + Parameters + ----------- + query: :class:`str` + The string that the username's start with. An empty string + requests all members. + limit: :class:`int` + The maximum number of members to send back. This must be + a number between 1 and 1000. + cache: :class:`bool` + Whether to cache the members internally. This makes operations + such as :meth:`get_member` work for those that matched. + + Raises + ------- + asyncio.TimeoutError + The query timed out waiting for the members. + + Returns + -------- + List[:class:`Member`] + The list of members that have matched the query. + """ + limit = limit or 5 + return await self._state.query_members(self, query=query, limit=limit, cache=cache) diff --git a/discord/state.py b/discord/state.py index 913911548..1fefa8008 100644 --- a/discord/state.py +++ b/discord/state.py @@ -51,6 +51,7 @@ from .object import Object class ListenerType(Enum): chunk = 0 + query_members = 1 Listener = namedtuple('Listener', ('type', 'future', 'predicate')) log = logging.getLogger(__name__) @@ -293,6 +294,32 @@ class ConnectionState: except asyncio.TimeoutError: log.info('Somehow timed out waiting for chunks.') + async def query_members(self, guild, query, limit, cache): + guild_id = guild.id + ws = self._get_websocket(guild_id) + if ws is None: + raise RuntimeError('Somehow do not have a websocket for this guild_id') + + # Limits over 1000 cannot be supported since + # the main use case for this is guild_subscriptions being disabled + # and they don't receive GUILD_MEMBER events which make computing + # member_count impossible. The only way to fix it is by limiting + # the limit parameter to 1 to 1000. + future = self.receive_member_query(guild_id, query) + try: + # start the query operation + await ws.request_chunks(guild_id, query, limit) + members = await asyncio.wait_for(future, timeout=5.0, loop=self.loop) + + if cache: + for member in members: + guild._add_member(member) + + return members + except asyncio.TimeoutError: + log.info('Timed out waiting for chunks with query %r and limit %d for guild_id %d', query, limit, guild_id) + raise + async def _delay_ready(self): try: launch = self._ready_state.launch @@ -792,15 +819,15 @@ class ConnectionState: def parse_guild_members_chunk(self, data): guild_id = int(data['guild_id']) guild = self._get_guild(guild_id) - members = data.get('members', []) - for member in members: - m = Member(guild=guild, data=member, state=self) - existing = guild.get_member(m.id) - if existing is None or existing.joined_at is None: - guild._add_member(m) - + members = [Member(guild=guild, data=member, state=self) for member in data.get('members', [])] log.info('Processed a chunk for %s members in guild ID %s.', len(members), guild_id) + if self._cache_members: + for member in members: + guild._add_member(member) + self.process_listeners(ListenerType.chunk, guild, len(members)) + names = [x.name.lower() for x in members] + self.process_listeners(ListenerType.query_members, (guild_id, names), members) def parse_guild_integrations_update(self, data): guild = self._get_guild(int(data['guild_id'])) @@ -930,6 +957,15 @@ class ConnectionState: self._listeners.append(listener) return future + def receive_member_query(self, guild_id, query): + def predicate(args, *, guild_id=guild_id, query=query.lower()): + request_guild_id, names = args + return request_guild_id == guild_id and all(n.startswith(query) for n in names) + future = self.loop.create_future() + listener = Listener(ListenerType.query_members, future, predicate) + self._listeners.append(listener) + return future + class AutoShardedConnectionState(ConnectionState): def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) diff --git a/discord/utils.py b/discord/utils.py index 48942a13d..490ab7678 100644 --- a/discord/utils.py +++ b/discord/utils.py @@ -323,11 +323,13 @@ async def async_all(gen, *, check=_isawaitable): return True async def sane_wait_for(futures, *, timeout, loop): - _, pending = await asyncio.wait(futures, timeout=timeout, loop=loop) + done, pending = await asyncio.wait(futures, timeout=timeout, return_when=asyncio.ALL_COMPLETED, loop=loop) if len(pending) != 0: raise asyncio.TimeoutError() + return done + def valid_icon_size(size): """Icons must be power of 2 within [16, 4096].""" return not size & (size - 1) and size in range(16, 4097)