From a1c618215e7c3222e07093b50f166bcaf53dcfc8 Mon Sep 17 00:00:00 2001 From: Josh Date: Mon, 14 Mar 2022 11:03:45 +1000 Subject: [PATCH] [commads] Change cog/extension load/unload methods to be async --- discord/__main__.py | 14 +++-- discord/ext/commands/bot.py | 87 ++++++++++++++++++---------- discord/ext/commands/cog.py | 36 +++++++++--- discord/ext/commands/core.py | 1 - docs/ext/commands/cogs.rst | 5 +- docs/ext/commands/extensions.rst | 10 ++-- docs/migrating.rst | 99 ++++++++++++++++++++++++++++++++ examples/basic_voice.py | 8 ++- 8 files changed, 210 insertions(+), 50 deletions(-) diff --git a/discord/__main__.py b/discord/__main__.py index 272750079..570dcbb86 100644 --- a/discord/__main__.py +++ b/discord/__main__.py @@ -66,9 +66,11 @@ import config class Bot(commands.{base}): def __init__(self, **kwargs): super().__init__(command_prefix=commands.when_mentioned_or('{prefix}'), **kwargs) + + async def setup_hook(self): for cog in config.cogs: try: - self.load_extension(cog) + await self.load_extension(cog) except Exception as exc: print(f'Could not load extension {{cog}} due to {{exc.__class__.__name__}}: {{exc}}') @@ -122,12 +124,16 @@ class {name}(commands.Cog{attrs}): def __init__(self, bot): self.bot = bot {extra} -def setup(bot): - bot.add_cog({name}(bot)) +async def setup(bot): + await bot.add_cog({name}(bot)) ''' _cog_extras = ''' - def cog_unload(self): + async def cog_load(self): + # loading logic goes here + pass + + async def cog_unload(self): # clean up logic goes here pass diff --git a/discord/ext/commands/bot.py b/discord/ext/commands/bot.py index f13643675..1ee2709ce 100644 --- a/discord/ext/commands/bot.py +++ b/discord/ext/commands/bot.py @@ -170,13 +170,13 @@ class BotBase(GroupMixin): async def close(self) -> None: for extension in tuple(self.__extensions): try: - self.unload_extension(extension) + await self.unload_extension(extension) except Exception: pass for cog in tuple(self.__cogs): try: - self.remove_cog(cog) + await self.remove_cog(cog) except Exception: pass @@ -528,7 +528,7 @@ class BotBase(GroupMixin): # cogs - def add_cog( + async def add_cog( self, cog: Cog, /, @@ -537,13 +537,20 @@ class BotBase(GroupMixin): guild: Optional[Snowflake] = MISSING, guilds: List[Snowflake] = MISSING, ) -> None: - """Adds a "cog" to the bot. + """|coro| + + Adds a "cog" to the bot. A cog is a class that has its own event listeners and commands. If the cog is a :class:`.app_commands.Group` then it is added to the bot's :class:`~discord.app_commands.CommandTree` as well. + .. note:: + + Exceptions raised inside a `class`:.Cog:'s :meth:`~.Cog.cog_load` method will be + propagated to the caller. + .. versionchanged:: 2.0 :exc:`.ClientException` is raised when a cog with the same name @@ -553,6 +560,10 @@ class BotBase(GroupMixin): ``cog`` parameter is now positional-only. + .. versionchanged:: 2.0 + + This method is now a :term:`coroutine`. + Parameters ----------- cog: :class:`.Cog` @@ -595,12 +606,12 @@ class BotBase(GroupMixin): if existing is not None: if not override: raise discord.ClientException(f'Cog named {cog_name!r} already loaded') - self.remove_cog(cog_name, guild=guild, guilds=guilds) + await self.remove_cog(cog_name, guild=guild, guilds=guilds) if isinstance(cog, app_commands.Group): self.__tree.add_command(cog, override=override, guild=guild, guilds=guilds) - cog = cog._inject(self, override=override, guild=guild, guilds=guilds) + cog = await cog._inject(self, override=override, guild=guild, guilds=guilds) self.__cogs[cog_name] = cog def get_cog(self, name: str, /) -> Optional[Cog]: @@ -626,14 +637,16 @@ class BotBase(GroupMixin): """ return self.__cogs.get(name) - def remove_cog( + async def remove_cog( self, name: str, /, guild: Optional[Snowflake] = MISSING, guilds: List[Snowflake] = MISSING, ) -> Optional[Cog]: - """Removes a cog from the bot and returns it. + """|coro| + + Removes a cog from the bot and returns it. All registered commands and event listeners that the cog has registered will be removed as well. @@ -644,6 +657,10 @@ class BotBase(GroupMixin): ``name`` parameter is now positional-only. + .. versionchanged:: 2.0 + + This method is now a :term:`coroutine`. + Parameters ----------- name: :class:`str` @@ -684,7 +701,7 @@ class BotBase(GroupMixin): for guild_id in guild_ids: self.__tree.remove_command(name, guild=discord.Object(guild_id)) - cog._eject(self, guild_ids=guild_ids) + await cog._eject(self, guild_ids=guild_ids) return cog @@ -695,12 +712,12 @@ class BotBase(GroupMixin): # extensions - def _remove_module_references(self, name: str) -> None: + async def _remove_module_references(self, name: str) -> None: # find all references to the module # remove the cogs registered from the module for cogname, cog in self.__cogs.copy().items(): if _is_submodule(name, cog.__module__): - self.remove_cog(cogname) + await self.remove_cog(cogname) # remove all the commands from the module for cmd in self.all_commands.copy().values(): @@ -722,14 +739,14 @@ class BotBase(GroupMixin): # remove all relevant application commands from the tree self.__tree._remove_with_module(name) - def _call_module_finalizers(self, lib: types.ModuleType, key: str) -> None: + async def _call_module_finalizers(self, lib: types.ModuleType, key: str) -> None: try: func = getattr(lib, 'teardown') except AttributeError: pass else: try: - func(self) + await func(self) except Exception: pass finally: @@ -740,7 +757,7 @@ class BotBase(GroupMixin): if _is_submodule(name, module): del sys.modules[module] - def _load_from_module_spec(self, spec: importlib.machinery.ModuleSpec, key: str) -> None: + async def _load_from_module_spec(self, spec: importlib.machinery.ModuleSpec, key: str) -> None: # precondition: key not in self.__extensions lib = importlib.util.module_from_spec(spec) sys.modules[key] = lib @@ -757,11 +774,11 @@ class BotBase(GroupMixin): raise errors.NoEntryPointError(key) try: - setup(self) + await setup(self) except Exception as e: del sys.modules[key] - self._remove_module_references(lib.__name__) - self._call_module_finalizers(lib, key) + await self._remove_module_references(lib.__name__) + await self._call_module_finalizers(lib, key) raise errors.ExtensionFailed(key, e) from e else: self.__extensions[key] = lib @@ -772,8 +789,10 @@ class BotBase(GroupMixin): except ImportError: raise errors.ExtensionNotFound(name) - def load_extension(self, name: str, *, package: Optional[str] = None) -> None: - """Loads an extension. + async def load_extension(self, name: str, *, package: Optional[str] = None) -> None: + """|coro| + + Loads an extension. An extension is a python module that contains commands, cogs, or listeners. @@ -782,6 +801,10 @@ class BotBase(GroupMixin): the entry point on what to do when the extension is loaded. This entry point must have a single argument, the ``bot``. + .. versionchanged:: 2.0 + + This method is now a :term:`coroutine`. + Parameters ------------ name: :class:`str` @@ -817,10 +840,12 @@ class BotBase(GroupMixin): if spec is None: raise errors.ExtensionNotFound(name) - self._load_from_module_spec(spec, name) + await self._load_from_module_spec(spec, name) - def unload_extension(self, name: str, *, package: Optional[str] = None) -> None: - """Unloads an extension. + async def unload_extension(self, name: str, *, package: Optional[str] = None) -> None: + """|coro| + + Unloads an extension. When the extension is unloaded, all commands, listeners, and cogs are removed from the bot and the module is un-imported. @@ -830,6 +855,10 @@ class BotBase(GroupMixin): parameter, the ``bot``, similar to ``setup`` from :meth:`~.Bot.load_extension`. + .. versionchanged:: 2.0 + + This method is now a :term:`coroutine`. + Parameters ------------ name: :class:`str` @@ -857,10 +886,10 @@ class BotBase(GroupMixin): if lib is None: raise errors.ExtensionNotLoaded(name) - self._remove_module_references(lib.__name__) - self._call_module_finalizers(lib, name) + await self._remove_module_references(lib.__name__) + await self._call_module_finalizers(lib, name) - def reload_extension(self, name: str, *, package: Optional[str] = None) -> None: + async def reload_extension(self, name: str, *, package: Optional[str] = None) -> None: """Atomically reloads an extension. This replaces the extension with the same extension, only refreshed. This is @@ -911,14 +940,14 @@ class BotBase(GroupMixin): try: # Unload and then load the module... - self._remove_module_references(lib.__name__) - self._call_module_finalizers(lib, name) - self.load_extension(name) + await self._remove_module_references(lib.__name__) + await self._call_module_finalizers(lib, name) + await self.load_extension(name) except Exception: # if the load failed, the remnants should have been # cleaned from the load_extension function call # so let's load it from our old compiled library. - lib.setup(self) # type: ignore + await lib.setup(self) # type: ignore self.__extensions[name] = lib # revert sys.modules back to normal and raise back to caller diff --git a/discord/ext/commands/cog.py b/discord/ext/commands/cog.py index 1ffbbe16d..d16584a39 100644 --- a/discord/ext/commands/cog.py +++ b/discord/ext/commands/cog.py @@ -26,6 +26,7 @@ from __future__ import annotations import inspect import discord from discord import app_commands +from discord.utils import maybe_coroutine from typing import Any, Callable, Dict, Generator, Iterable, List, Optional, TYPE_CHECKING, Tuple, TypeVar, Union @@ -377,13 +378,30 @@ class Cog(metaclass=CogMeta): return not hasattr(self.cog_command_error.__func__, '__cog_special_method__') @_cog_special_method - def cog_unload(self) -> None: - """A special method that is called when the cog gets removed. + async def cog_load(self) -> None: + """|maybecoro| - This function **cannot** be a coroutine. It must be a regular - function. + A special method that is called when the cog gets loaded. + + Subclasses must replace this if they want special asynchronous loading behaviour. + Note that the ``__init__`` special method does not allow asynchronous code to run + inside it, thus this is helpful for setting up code that needs to be asynchronous. + + .. versionadded:: 2.0 + """ + pass + + @_cog_special_method + async def cog_unload(self) -> None: + """|maybecoro| + + A special method that is called when the cog gets removed. Subclasses must replace this if they want special unloading behaviour. + + .. versionchanged:: 2.0 + + This method can now be a :term:`coroutine`. """ pass @@ -466,9 +484,13 @@ class Cog(metaclass=CogMeta): """ pass - def _inject(self, bot: BotBase, override: bool, guild: Optional[Snowflake], guilds: List[Snowflake]) -> Self: + async def _inject(self, bot: BotBase, override: bool, guild: Optional[Snowflake], guilds: List[Snowflake]) -> Self: cls = self.__class__ + # we'll call this first so that errors can propagate without + # having to worry about undoing anything + await maybe_coroutine(self.cog_load) + # realistically, the only thing that can cause loading errors # is essentially just the command loading, which raises if there are # duplicates. When this condition is met, we want to undo all what @@ -507,7 +529,7 @@ class Cog(metaclass=CogMeta): return self - def _eject(self, bot: BotBase, guild_ids: Optional[Iterable[int]]) -> None: + async def _eject(self, bot: BotBase, guild_ids: Optional[Iterable[int]]) -> None: cls = self.__class__ try: @@ -534,6 +556,6 @@ class Cog(metaclass=CogMeta): bot.remove_check(self.bot_check_once, call_once=True) finally: try: - self.cog_unload() + await maybe_coroutine(self.cog_unload) except Exception: pass diff --git a/discord/ext/commands/core.py b/discord/ext/commands/core.py index 741280d36..57eeba70b 100644 --- a/discord/ext/commands/core.py +++ b/discord/ext/commands/core.py @@ -2350,7 +2350,6 @@ def before_invoke(coro) -> Callable[[T], T]: async def why(self, ctx): # Output: await ctx.send('because someone made me') - bot.add_cog(What()) """ def decorator(func: Union[Command, CoroFunc]) -> Union[Command, CoroFunc]: diff --git a/docs/ext/commands/cogs.rst b/docs/ext/commands/cogs.rst index 25bb1a0b8..df2d9cdc1 100644 --- a/docs/ext/commands/cogs.rst +++ b/docs/ext/commands/cogs.rst @@ -58,7 +58,7 @@ Once you have defined your cogs, you need to tell the bot to register the cogs t .. code-block:: python3 - bot.add_cog(Greetings(bot)) + await bot.add_cog(Greetings(bot)) This binds the cog to the bot, adding all commands and listeners to the bot automatically. @@ -66,7 +66,7 @@ Note that we reference the cog by name, which we can override through :ref:`ext_ .. code-block:: python3 - bot.remove_cog('Greetings') + await bot.remove_cog('Greetings') Using Cogs ------------- @@ -112,6 +112,7 @@ As cogs get more complicated and have more commands, there comes a point where w They are as follows: +- :meth:`.Cog.cog_load` - :meth:`.Cog.cog_unload` - :meth:`.Cog.cog_check` - :meth:`.Cog.cog_command_error` diff --git a/docs/ext/commands/extensions.rst b/docs/ext/commands/extensions.rst index 20aa6e128..136e8c920 100644 --- a/docs/ext/commands/extensions.rst +++ b/docs/ext/commands/extensions.rst @@ -24,10 +24,10 @@ An example extension looks like this: async def hello(ctx): await ctx.send(f'Hello {ctx.author.display_name}.') - def setup(bot): + async def setup(bot): bot.add_command(hello) -In this example we define a simple command, and when the extension is loaded this command is added to the bot. Now the final step to this is loading the extension, which we do by calling :meth:`.Bot.load_extension`. To load this extension we call ``bot.load_extension('hello')``. +In this example we define a simple command, and when the extension is loaded this command is added to the bot. Now the final step to this is loading the extension, which we do by calling :meth:`.Bot.load_extension`. To load this extension we call ``await bot.load_extension('hello')``. .. admonition:: Cogs :class: helpful @@ -45,7 +45,7 @@ When you make a change to the extension and want to reload the references, the l .. code-block:: python3 - >>> bot.reload_extension('hello') + >>> await bot.reload_extension('hello') Once the extension reloads, any changes that we did will be applied. This is useful if we want to add or remove functionality without restarting our bot. If an error occurred during the reloading process, the bot will pretend as if the reload never happened. @@ -57,8 +57,8 @@ Although rare, sometimes an extension needs to clean-up or know when it's being .. code-block:: python3 :caption: basic_ext.py - def setup(bot): + async def setup(bot): print('I am being loaded!') - def teardown(bot): + async def teardown(bot): print('I am being unloaded!') diff --git a/docs/migrating.rst b/docs/migrating.rst index 7f2e2e677..38a5dd265 100644 --- a/docs/migrating.rst +++ b/docs/migrating.rst @@ -65,6 +65,54 @@ The following have been removed: - ``User.mutual_friends`` attribute - ``User.relationship`` attribute +.. _migrating_2_0_client_async_setup: + +asyncio Event Loop Changes +--------------------------- + +Python 3.7 introduced a new helper function :func:`asyncio.run` which automatically creates and destroys the asynchronous event loop. + +In order to support this, the way discord.py handles the :mod:`asyncio` event loop has changed. + +This allows you to rather than using :meth:`Client.run` create your own asynchronous loop to setup other asynchronous code as needed. + +Quick example: + +.. code-block:: python + + client = discord.Client() + + async def main(): + # do other async things + await my_async_function() + + # start the client + async with client: + await client.start(TOKEN) + + asyncio.run(main()) + +A new :meth:`~Client.setup_hook` method has also been added to the :class:`Client` class. +This method is called after login but before connecting to the discord gateway. + +It is intended to be used to setup various bot features in an asynchronous context. + +:meth:`~Client.setup_hook` can be defined by subclassing the :class:`Client` class. + +Quick example: + +.. code-block:: python + + class MyClient(discord.Client): + async def setup_hook(self): + print('This is asynchronous!') + + client = MyClient() + client.run(TOKEN) + +In parallel with this change, changes were made to loading and unloading of commands extension extensions and cogs, +see :ref:`migrating_2_0_commands_extension_cog_async` for more information. + Abstract Base Classes Changes ------------------------------- @@ -1095,6 +1143,55 @@ The following changes have been made: Command Extension Changes --------------------------- +.. _migrating_2_0_commands_extension_cog_async: + +Extension and Cog Loading / Unloading is Now Asynchronous +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +As an extension to the :ref:`asyncio changes ` the loading and unloading of extensions and cogs is now asynchronous. + +To accommodate this, the following changes have been made: + +- the ``setup`` and ``teardown`` functions in extensions must now be coroutines. +- :meth:`ext.commands.Bot.load_extension` must now be awaited. +- :meth:`ext.commands.Bot.unload_extension` must now be awaited. +- :meth:`ext.commands.Bot.reload_extension` must now be awaited. +- :meth:`ext.commands.Bot.add_cog` must now be awaited. +- :meth:`ext.commands.Bot.remove_cog` must now be awaited. + +Quick example of an extension setup function: + +.. code:: python + + # before + def setup(bot): + bot.add_cog(MyCog(bot)) + + #after + async def setup(bot): + await bot.add_cog(MyCog(bot)) + +Quick example of loading an extension: + +.. code:: python + + #before + bot.load_extension('my_extension') + + #after using setup_hook + class MyBot(commands.Bot): + async def setup_hook(self): + await self.load_extension('my_extension') + + # after using async_with + async def main(): + async with bot: + await bot.load_extension('my_extension') + await bot.start(TOKEN) + + asyncio.run(main()) + + Converters Are Now Generic Runtime Protocols ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ @@ -1192,6 +1289,8 @@ Miscellanous Changes - :attr:`ext.commands.Context.channel` may now be a :class:`PartialMessageable`. - ``MissingPermissions.missing_perms`` has been renamed to :attr:`ext.commands.MissingPermissions.missing_permissions`. - ``BotMissingPermissions.missing_perms`` has been renamed to :attr:`ext.commands.BotMissingPermissions.missing_permissions`. +- :meth:`ext.commands.Cog.cog_load` has been added as part of the :ref:`migrating_2_0_commands_extension_cog_async` changes. +- :meth:`ext.commands.Cog.cog_unload` may now be a :term:`coroutine` due to the :ref:`migrating_2_0_commands_extension_cog_async` changes. .. _migrating_2_0_tasks: diff --git a/examples/basic_voice.py b/examples/basic_voice.py index bae1174a5..273fa202e 100644 --- a/examples/basic_voice.py +++ b/examples/basic_voice.py @@ -137,5 +137,9 @@ async def on_ready(): print(f'Logged in as {bot.user} (ID: {bot.user.id})') print('------') -bot.add_cog(Music(bot)) -bot.run('token') +async def main(): + async with bot: + await bot.add_cog(Music(bot)) + await bot.start('token') + +asyncio.run(main())