diff --git a/discord/__init__.py b/discord/__init__.py index 5e07acecf..63d6a849a 100644 --- a/discord/__init__.py +++ b/discord/__init__.py @@ -49,7 +49,7 @@ from . import utils, opus, abc from .enums import * from .embeds import Embed from .mentions import AllowedMentions -from .shard import AutoShardedClient +from .shard import AutoShardedClient, ShardInfo from .player import * from .webhook import * from .voice_client import VoiceClient diff --git a/discord/shard.py b/discord/shard.py index e3835187b..36d7477e6 100644 --- a/discord/shard.py +++ b/discord/shard.py @@ -103,6 +103,10 @@ class Shard: self._cancel_task() await self.ws.close(code=1000) + async def disconnect(self): + await self.close() + self._dispatch('shard_disconnect', self.id) + async def _handle_disconnect(self, e): self._dispatch('disconnect') self._dispatch('shard_disconnect', self.id) @@ -178,6 +182,70 @@ class Shard: else: self.launch() +class ShardInfo: + """A class that gives information and control over a specific shard. + + You can retrieve this object via :meth:`AutoShardedClient.get_shard` + or :attr:`AutoShardedClient.shards`. + + .. versionadded:: 1.4 + + Attributes + ------------ + id: :class:`int` + The shard ID for this shard. + shard_count: Optional[:class:`int`] + The shard count for this cluster. If this is ``None`` then the bot has not started yet. + """ + + __slots__ = ('_parent', 'id', 'shard_count') + + def __init__(self, parent, shard_count): + self._parent = parent + self.id = parent.id + self.shard_count = shard_count + + def is_closed(self): + """:class:`bool`: Whether the shard connection is currently closed.""" + return not self._parent.ws.open + + async def disconnect(self): + """|coro| + + Disconnects a shard. When this is called, the shard connection will no + longer be open. + + If the shard is already disconnected this does nothing. + """ + if self.is_closed(): + return + + await self._parent.disconnect() + + async def reconnect(self): + """|coro| + + Disconnects and then connects the shard again. + """ + if not self.is_closed(): + await self._parent.disconnect() + await self._parent.reconnect() + + async def connect(self): + """|coro| + + Connects a shard. If the shard is already connected this does nothing. + """ + if not self.is_closed(): + return + + await self._parent.reconnect() + + @property + def latency(self): + """:class:`float`: Measures latency between a HEARTBEAT and a HEARTBEAT_ACK in seconds for this shard.""" + return self._parent.ws.latency + class AutoShardedClient(Client): """A client similar to :class:`Client` except it handles the complications of sharding for the user into a more manageable and transparent single @@ -221,14 +289,14 @@ class AutoShardedClient(Client): # instead of a single websocket, we have multiple # the key is the shard_id - self.shards = {} + self.__shards = {} self._connection._get_websocket = self._get_websocket - self._queue = asyncio.PriorityQueue() + self.__queue = asyncio.PriorityQueue() def _get_websocket(self, guild_id=None, *, shard_id=None): if shard_id is None: shard_id = (guild_id >> 22) % self.shard_count - return self.shards[shard_id].ws + return self.__shards[shard_id].ws @property def latency(self): @@ -238,9 +306,9 @@ class AutoShardedClient(Client): latency of every shard's latency. To get a list of shard latency, check the :attr:`latencies` property. Returns ``nan`` if there are no shards ready. """ - if not self.shards: + if not self.__shards: return float('nan') - return sum(latency for _, latency in self.latencies) / len(self.shards) + return sum(latency for _, latency in self.latencies) / len(self.__shards) @property def latencies(self): @@ -248,7 +316,21 @@ class AutoShardedClient(Client): This returns a list of tuples with elements ``(shard_id, latency)``. """ - return [(shard_id, shard.ws.latency) for shard_id, shard in self.shards.items()] + return [(shard_id, shard.ws.latency) for shard_id, shard in self.__shards.items()] + + def get_shard(self, shard_id): + """Optional[:class:`ShardInfo`]: Gets the shard information at a given shard ID or ``None`` if not found.""" + try: + parent = self.__shards[shard_id] + except KeyError: + return None + else: + return ShardInfo(parent, self.shard_count) + + @utils.cached_property + def shards(self): + """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() } async def request_offline_members(self, *guilds): r"""|coro| @@ -291,7 +373,7 @@ class AutoShardedClient(Client): return await self.launch_shard(gateway, shard_id) # keep reading the shard while others connect - self.shards[shard_id] = ret = Shard(ws, self) + self.__shards[shard_id] = ret = Shard(ws, self) ret.launch() async def launch_shards(self): @@ -316,7 +398,7 @@ class AutoShardedClient(Client): await self.launch_shards() while not self.is_closed(): - item = await self._queue.get() + item = await self.__queue.get() if item.type == EventType.close: await self.close() if isinstance(item.error, ConnectionClosed) and item.error.code != 1000: @@ -346,7 +428,7 @@ class AutoShardedClient(Client): except Exception: pass - to_close = [asyncio.ensure_future(shard.close(), loop=self.loop) for shard in self.shards.values()] + to_close = [asyncio.ensure_future(shard.close(), loop=self.loop) for shard in self.__shards.values()] if to_close: await asyncio.wait(to_close) @@ -395,12 +477,12 @@ class AutoShardedClient(Client): status = str(status) if shard_id is None: - for shard in self.shards.values(): + for shard in self.__shards.values(): await shard.ws.change_presence(activity=activity, status=status, afk=afk) guilds = self._connection.guilds else: - shard = self.shards[shard_id] + shard = self.__shards[shard_id] await shard.ws.change_presence(activity=activity, status=status, afk=afk) guilds = [g for g in self._connection.guilds if g.shard_id == shard_id] diff --git a/docs/api.rst b/docs/api.rst index 601a6c012..8fdd32c41 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -2756,7 +2756,7 @@ Data Classes Some classes are just there to be data containers, this lists them. Unlike :ref:`models ` you are allowed to create -these yourself, even if they can also be used to hold attributes. +most of these yourself, even if they can also be used to hold attributes. Nearly all classes here have :ref:`py:slots` defined which means that it is impossible to have dynamic attributes to the data classes. @@ -2837,22 +2837,28 @@ PermissionOverwrite .. autoclass:: PermissionOverwrite :members: +ShardInfo +~~~~~~~~~~~ + +.. autoclass:: ShardInfo() + :members: + SystemChannelFlags ~~~~~~~~~~~~~~~~~~~~ -.. autoclass:: SystemChannelFlags +.. autoclass:: SystemChannelFlags() :members: MessageFlags ~~~~~~~~~~~~ -.. autoclass:: MessageFlags +.. autoclass:: MessageFlags() :members: PublicUserFlags ~~~~~~~~~~~~~~~ -.. autoclass:: PublicUserFlags +.. autoclass:: PublicUserFlags() :members: