Add Guild.query_members to fetch members from the gateway.

4 files changed, 97 insertions, 8 deletions

diff --git a/discord/gateway.py b/discord/gateway.py
index 7d71425b..6848f120 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 4c87e9f0..5a95f5da 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 91391154..1fefa800 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 48942a13..490ab767 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)