From 9c61e10a55d5b776a8731eb9a6457dfbf3cc9e9f Mon Sep 17 00:00:00 2001 From: Rapptz Date: Sun, 13 Mar 2022 08:26:39 -0400 Subject: [PATCH] Move all async object creation to a proper initialisation point This should make it so no object is created with another loop --- discord/client.py | 54 +++++++++++++++++++++++++++-------------------- discord/http.py | 6 ++++-- discord/shard.py | 5 ++++- discord/state.py | 7 +++++- 4 files changed, 45 insertions(+), 27 deletions(-) diff --git a/discord/client.py b/discord/client.py index 6fb3faba6..fc3bea864 100644 --- a/discord/client.py +++ b/discord/client.py @@ -208,7 +208,7 @@ class Client: self._connection: ConnectionState = self._get_state(**options) self._connection.shard_count = self.shard_count self._closed: bool = False - self._ready: asyncio.Event = asyncio.Event() + self._ready: asyncio.Event = MISSING self._connection._get_websocket = self._get_websocket self._connection._get_client = lambda: self @@ -333,7 +333,7 @@ class Client: def is_ready(self) -> bool: """:class:`bool`: Specifies if the client's internal cache is ready for use.""" - return self._ready.is_set() + return self._ready is not MISSING and self._ready.is_set() async def _run_event( self, @@ -445,6 +445,32 @@ class Client: if not initial: await asyncio.sleep(5.0) + async def _async_setup_hook(self) -> None: + # Called whenever the client needs to initialise asyncio objects with a running loop + loop = asyncio.get_running_loop() + self.loop = loop + self.http.loop = loop + self._connection.loop = loop + await self._connection.async_setup() + + self._ready = asyncio.Event() + + async def setup_hook(self) -> None: + """|coro| + + A coroutine to be called to setup the bot, by default this is blank. + + To perform asynchronous setup after the bot is logged in but before + it has connected to the Websocket, overwrite this coroutine. + + This is only called once, in :meth:`login`, and will be called before + any events are dispatched, making it a better solution than doing such + setup in the :func:`~discord.on_ready` event. + + .. versionadded:: 2.0 + """ + pass + # login state management async def login(self, token: str) -> None: @@ -472,10 +498,7 @@ class Client: _log.info('logging in using static token') - loop = asyncio.get_running_loop() - self.loop = loop - self.http.loop = loop - self._connection.loop = loop + await self._async_setup_hook() data = await self.http.static_login(token.strip()) self._connection.user = ClientUser(state=self._connection, data=data) @@ -617,22 +640,6 @@ class Client: await self.login(token) await self.connect(reconnect=reconnect) - async def setup_hook(self) -> None: - """|coro| - - A coroutine to be called to setup the bot, by default this is blank. - - To perform asynchronous setup after the bot is logged in but before - it has connected to the Websocket, overwrite this coroutine. - - This is only called once, in :meth:`login`, and will be called before - any events are dispatched, making it a better solution than doing such - setup in the :func:`~discord.on_ready` event. - - .. versionadded:: 2.0 - """ - pass - def run(self, *args: Any, **kwargs: Any) -> None: """A blocking call that abstracts away the event loop initialisation from you. @@ -925,7 +932,8 @@ class Client: Waits until the client's internal cache is all ready. """ - await self._ready.wait() + if self._ready is not MISSING: + await self._ready.wait() def wait_for( self, diff --git a/discord/http.py b/discord/http.py index f73f55ab5..645111918 100644 --- a/discord/http.py +++ b/discord/http.py @@ -343,8 +343,7 @@ class HTTPClient: self.connector: aiohttp.BaseConnector = connector or MISSING self.__session: aiohttp.ClientSession = MISSING # filled in static_login self._locks: weakref.WeakValueDictionary = weakref.WeakValueDictionary() - self._global_over: asyncio.Event = asyncio.Event() - self._global_over.set() + self._global_over: asyncio.Event = MISSING self.token: Optional[str] = None self.bot_token: bool = False self.proxy: Optional[str] = proxy @@ -550,6 +549,9 @@ class HTTPClient: self.__session = aiohttp.ClientSession( connector=self.connector, ws_response_class=DiscordClientWebSocketResponse, loop=self.loop ) + self._global_over = asyncio.Event() + self._global_over.set() + old_token = self.token self.token = token diff --git a/discord/shard.py b/discord/shard.py index 417742acd..59a350515 100644 --- a/discord/shard.py +++ b/discord/shard.py @@ -424,8 +424,11 @@ class AutoShardedClient(Client): self._connection.shards_launched.set() - async def connect(self, *, reconnect: bool = True) -> None: + async def _async_setup_hook(self) -> None: + await super()._async_setup_hook() self.__queue = asyncio.PriorityQueue() + + async def connect(self, *, reconnect: bool = True) -> None: self._reconnect = reconnect await self.launch_shards() diff --git a/discord/state.py b/discord/state.py index 9e0569524..aa4e36eaf 100644 --- a/discord/state.py +++ b/discord/state.py @@ -300,6 +300,9 @@ class ConnectionState: else: await coro(*args, **kwargs) + async def async_setup(self) -> None: + pass + @property def self_id(self) -> Optional[int]: u = self.user @@ -1485,7 +1488,6 @@ class AutoShardedConnectionState(ConnectionState): def __init__(self, *args: Any, **kwargs: Any) -> None: super().__init__(*args, **kwargs) self.shard_ids: Union[List[int], range] = [] - self.shards_launched: asyncio.Event = asyncio.Event() def _update_message_references(self) -> None: # self._messages won't be None when this is called @@ -1500,6 +1502,9 @@ class AutoShardedConnectionState(ConnectionState): # channel will either be a TextChannel, Thread or Object msg._rebind_cached_references(new_guild, channel) # type: ignore + async def async_setup(self) -> None: + self.shards_launched: asyncio.Event = asyncio.Event() + async def chunker( self, guild_id: int,