diff --git a/discord/experiment.py b/discord/experiment.py index 556c4033c..ed83474cd 100644 --- a/discord/experiment.py +++ b/discord/experiment.py @@ -702,13 +702,13 @@ class UserExperiment: Returns the experiment's hash. - .. versionadded:: 2.1 - .. note:: In contrast to the wide range of data provided for guild experiments, user experiments do not reveal detailed rollout information, providing only the assigned bucket. + .. versionadded:: 2.1 + Attributes ---------- hash: :class:`int` diff --git a/discord/gateway.py b/discord/gateway.py index 5822b4a90..cc3872d7d 100644 --- a/discord/gateway.py +++ b/discord/gateway.py @@ -294,23 +294,24 @@ class DiscordWebSocket: _zlib_enabled: bool # fmt: off - DEFAULT_GATEWAY = yarl.URL('wss://gateway.discord.gg/') - DISPATCH = 0 - HEARTBEAT = 1 - IDENTIFY = 2 - PRESENCE = 3 - VOICE_STATE = 4 - VOICE_PING = 5 - RESUME = 6 - RECONNECT = 7 - REQUEST_MEMBERS = 8 - INVALIDATE_SESSION = 9 - HELLO = 10 - HEARTBEAT_ACK = 11 - GUILD_SYNC = 12 # :( - CALL_CONNECT = 13 - GUILD_SUBSCRIBE = 14 - REQUEST_COMMANDS = 24 + DEFAULT_GATEWAY = yarl.URL('wss://gateway.discord.gg/') + DISPATCH = 0 + HEARTBEAT = 1 + IDENTIFY = 2 + PRESENCE = 3 + VOICE_STATE = 4 + VOICE_PING = 5 + RESUME = 6 + RECONNECT = 7 + REQUEST_MEMBERS = 8 + INVALIDATE_SESSION = 9 + HELLO = 10 + HEARTBEAT_ACK = 11 + GUILD_SYNC = 12 # :( + CALL_CONNECT = 13 + GUILD_SUBSCRIBE = 14 + REQUEST_COMMANDS = 24 + SEARCH_RECENT_MEMBERS = 35 # fmt: on def __init__(self, socket: aiohttp.ClientWebSocketResponse, *, loop: asyncio.AbstractEventLoop) -> None: @@ -325,7 +326,7 @@ class DiscordWebSocket: self._keep_alive: Optional[KeepAliveHandler] = None self.thread_id: int = threading.get_ident() - # ws related stuff + # WS related stuff self.session_id: Optional[str] = None self.sequence: Optional[int] = None self._zlib: zlib._Decompress = zlib.decompressobj() @@ -855,6 +856,22 @@ class DiscordWebSocket: await self.send_as_json(payload) + async def search_recent_members( + self, guild_id: Snowflake, query: str = '', *, after: Optional[Snowflake] = None, nonce: Optional[str] = None + ) -> None: + payload = { + 'op': self.SEARCH_RECENT_MEMBERS, + 'd': { + 'guild_id': str(guild_id), + 'query': query, + 'continuation_token': str(after) if after else None, + }, + } + if nonce is not None: + payload['d']['nonce'] = nonce + + await self.send_as_json(payload) + async def close(self, code: int = 4000) -> None: if self._keep_alive: self._keep_alive.stop() diff --git a/discord/guild.py b/discord/guild.py index 3c2e1b604..ad9171691 100644 --- a/discord/guild.py +++ b/discord/guild.py @@ -4799,6 +4799,7 @@ class Guild(Hashable): cache: :class:`bool` Whether to cache the members internally. This makes operations such as :meth:`get_member` work for those that matched. + The cache will not be kept updated unless ``subscribe`` is set to ``True``. 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. @@ -4837,6 +4838,61 @@ class Guild(Hashable): await self._state.ws.request_lazy_guild(self.id, members=ids) return members + async def query_recent_members( + self, + query: Optional[str] = None, + *, + limit: int = 1000, + cache: bool = True, + subscribe: bool = False, + ) -> List[Member]: + """|coro| + + Request the most recent 10,000 joined members of this guild. + This is a websocket operation. + + .. note:: + + This operation does not return presences. + + .. versionadded:: 2.1 + + Parameters + ----------- + query: Optional[:class:`str`] + The string that the username or nickname should start with, if any. + limit: :class:`int` + The maximum number of members to send back. This must be + a number between 1 and 10,000. + cache: :class:`bool` + Whether to cache the members internally. This makes operations + such as :meth:`get_member` work for those that matched. + The cache will not be kept updated unless ``subscribe`` is set to ``True``. + subscribe: :class:`bool` + Whether to subscribe to the resulting members. This will keep their info and presence updated. + This requires another request, and defaults to ``False``. + + Raises + ------- + asyncio.TimeoutError + The query timed out waiting for the members. + TypeError + Invalid parameters were passed to the function. + + Returns + -------- + List[:class:`Member`] + The list of members that have matched the query. + """ + limit = min(10000, limit or 1) + members = await self._state.search_recent_members(self, query or '', limit, cache) + if subscribe: + ids: List[_Snowflake] = [str(m.id) for m in members] + for i in range(0, len(ids), 750): + subs = ids[i : i + 750] + await self._state.ws.request_lazy_guild(self.id, members=subs) + return members + async def change_voice_state( self, *, diff --git a/discord/state.py b/discord/state.py index c0e25c0df..8a25404ea 100644 --- a/discord/state.py +++ b/discord/state.py @@ -140,24 +140,55 @@ _log = logging.getLogger(__name__) class ChunkRequest: + __slots__ = ( + 'guild_id', + 'resolver', + 'loop', + 'limit', + 'remaining', + 'cache', + 'oneshot', + 'nonce', + 'buffer', + 'last_buffer', + 'waiters', + ) + def __init__( self, guild_id: int, loop: asyncio.AbstractEventLoop, resolver: Callable[[int], Any], *, + limit: Optional[int] = None, cache: bool = True, + oneshot: bool = True, ) -> None: self.guild_id: int = guild_id self.resolver: Callable[[int], Any] = resolver self.loop: asyncio.AbstractEventLoop = loop + self.limit: Optional[int] = limit + self.remaining: int = limit or 0 self.cache: bool = cache + self.oneshot: bool = oneshot self.nonce: str = str(utils.time_snowflake(utils.utcnow())) self.buffer: List[Member] = [] + self.last_buffer: Optional[List[Member]] = None self.waiters: List[asyncio.Future[List[Member]]] = [] def add_members(self, members: List[Member]) -> None: + unique_members = set(members) + if self.limit is not None: + if self.remaining <= 0: + return + + members = list(unique_members)[: self.remaining] + self.remaining -= len(unique_members) + else: + members = list(unique_members) + self.buffer.extend(members) + if self.cache: guild = self.resolver(self.guild_id) if guild is None: @@ -166,6 +197,9 @@ class ChunkRequest: for member in members: guild._add_member(member) + if not self.oneshot: + self.last_buffer = members + async def wait(self) -> List[Member]: future = self.loop.create_future() self.waiters.append(future) @@ -180,12 +214,28 @@ class ChunkRequest: return future def done(self) -> None: + result = self.buffer if self.oneshot else self.last_buffer or self.buffer for future in self.waiters: if not future.done(): - future.set_result(self.buffer) + future.set_result(result) class MemberSidebar: + __slots__ = ( + 'guild', + 'channels', + 'chunk', + 'delay', + 'cache', + 'loop', + 'safe_override', + 'ranges', + 'subscribing', + 'buffer', + 'exception', + 'waiters', + ) + def __init__( self, guild: Guild, @@ -301,6 +351,7 @@ class MemberSidebar: return list(ret) def add_members(self, members: List[Member]) -> None: + members = list(set(members)) self.buffer.extend(members) if self.cache: guild = self.guild @@ -653,7 +704,8 @@ class ConnectionState: request.add_members(members) if complete: request.done() - removed.append(key) + if request.oneshot: + removed.append(key) for key in removed: del self._chunk_requests[key] @@ -905,9 +957,18 @@ class ConnectionState: return self.ws.request_lazy_guild(guild_id, typing=typing, activities=activities, threads=threads) def chunker( - self, guild_id: int, query: str = '', limit: int = 0, presences: bool = True, *, nonce: Optional[str] = None + self, + guild_id: int, + query: Optional[str] = '', + limit: int = 0, + presences: bool = True, + *, + user_ids: Optional[List[Snowflake]] = None, + nonce: Optional[str] = None, ): - return self.ws.request_chunks([guild_id], query=query, limit=limit, presences=presences, nonce=nonce) + return self.ws.request_chunks( + [guild_id], query=query, limit=limit, presences=presences, user_ids=user_ids, nonce=nonce + ) async def query_members( self, @@ -923,14 +984,51 @@ class ConnectionState: self._chunk_requests[request.nonce] = request try: - await self.ws.request_chunks( - [guild_id], query=query, limit=limit, user_ids=user_ids, presences=presences, nonce=request.nonce + await self.chunker( + guild_id, query=query, limit=limit, presences=presences, user_ids=user_ids, nonce=request.nonce ) return await asyncio.wait_for(request.wait(), 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 + async def search_recent_members( + self, + guild: Guild, + query: str = '', + limit: Optional[int] = None, + cache: bool = False, + ) -> List[Member]: + guild_id = guild.id + request = ChunkRequest(guild.id, self.loop, self._get_guild, limit=limit, cache=cache, oneshot=False) + self._chunk_requests[request.nonce] = request + + # Unlike query members, this OP is paginated + old_continuation_token = None + continuation_token = None + while True: + try: + await self.ws.search_recent_members(guild_id, query=query, nonce=request.nonce, after=continuation_token) + returned = await asyncio.wait_for(request.wait(), timeout=30.0) + except asyncio.TimeoutError: + _log.warning( + 'Timed out waiting for search chunks with query %r and limit %d for guild_id %d.', query, limit, guild_id + ) + raise + + if (limit is not None and request.remaining < 1) or len(returned) < 1: + break + + # Sort the members by joined_at timestamp and grab the oldest one + request.buffer.sort(key=lambda m: m.joined_at or utils.utcnow()) + old_continuation_token = continuation_token + continuation_token = request.buffer[0].id + if continuation_token == old_continuation_token: + break + + self._chunk_requests.pop(request.nonce, None) + return list(set(request.buffer)) + async def _delay_ready(self) -> None: try: states = [] diff --git a/discord/user.py b/discord/user.py index c95e191e9..c5ce2fd50 100644 --- a/discord/user.py +++ b/discord/user.py @@ -821,7 +821,7 @@ class ClientUser(BaseUser): The hypesquad house you wish to change to. Could be ``None`` to leave the current house. username: :class:`str` - The new username you wish to change to. + The new username you wish to change to. discriminator: :class:`int` The new discriminator you wish to change to. This is a legacy concept that is no longer used. Can only be used if you have Nitro. @@ -1024,11 +1024,18 @@ class User(BaseUser, discord.abc.Connectable, discord.abc.Messageable): user['discriminator'], user.get('public_flags', 0), user.get('avatar_decoration'), - user.get('global_name') + user.get('global_name'), ) if original != modified: to_return = User._copy(self) - self.name, self._avatar, self.discriminator, self._public_flags, self._avatar_decoration, self.global_name = modified + ( + self.name, + self._avatar, + self.discriminator, + self._public_flags, + self._avatar_decoration, + self.global_name, + ) = modified # Signal to dispatch user_update return to_return, self