Merge pull request Pycord-Development#1411 from baronkobama/bridge-hooks

Add new hooks/handlers to bridge commands

discord/ext/bridge/core.py

@@ -23,6 +23,7 @@
 DEALINGS IN THE SOFTWARE.
 """
 from typing import Any, List, Union
+import asyncio
 
 import discord.commands.options
 from discord.commands import Option, SlashCommand
@@ -78,6 +79,9 @@ def __init__(self, callback, **kwargs):
         self.callback = callback
         self.kwargs = kwargs
 
+        self.ext_command = BridgeExtCommand(self.callback, **self.kwargs)
+        self.application_command = BridgeSlashCommand(self.callback, **self.kwargs)
+
     def get_ext_command(self):
         """A method to get the ext.commands version of this command.
 
@@ -86,8 +90,7 @@ def get_ext_command(self):
         :class:`BridgeExtCommand`
             The respective traditional (prefix-based) version of the command.
         """
-        command = BridgeExtCommand(self.callback, **self.kwargs)
-        return command
+        return self.ext_command
 
     def get_application_command(self):
         """A method to get the discord.commands version of this command.
@@ -97,8 +100,7 @@ def get_application_command(self):
         :class:`BridgeSlashCommand`
             The respective slash command version of the command.
         """
-        command = BridgeSlashCommand(self.callback, **self.kwargs)
-        return command
+        return self.application_command
 
     def add_to(self, bot: Union[ExtBot, ExtAutoShardedBot]) -> None:
         """Adds the command to a bot.
@@ -108,8 +110,95 @@ def add_to(self, bot: Union[ExtBot, ExtAutoShardedBot]) -> None:
         bot: Union[:class:`ExtBot`, :class:`ExtAutoShardedBot`]
             The bot to add the command to.
         """
-        bot.add_command(self.get_ext_command())
-        bot.add_application_command(self.get_application_command())
+
+        bot.add_command(self.ext_command)
+        bot.add_application_command(self.application_command)
+
+    def error(self, coro):
+        """A decorator that registers a coroutine as a local error handler.
+
+        This error handler is limited to the command it is defined to.
+        However, higher scope handlers (per-cog and global) are still
+        invoked afterwards as a catch-all. This handler also functions as
+        the handler for both the prefixed and slash versions of the command.
+
+        This error handler takes two parameters, a :class:`.BridgeContext` and
+        a :class:`~discord.DiscordException`.
+
+        Parameters
+        -----------
+        coro: :ref:`coroutine <coroutine>`
+            The coroutine to register as the local error handler.
+
+        Raises
+        -------
+        TypeError
+            The coroutine passed is not actually a coroutine.
+        """
+
+        if not asyncio.iscoroutinefunction(coro):
+            raise TypeError("The error handler must be a coroutine.")
+
+        self.ext_command.on_error = coro
+        self.application_command.on_error = coro
+
+        return coro
+
+    def before_invoke(self, coro):
+        """A decorator that registers a coroutine as a pre-invoke hook.
+
+        This hook is called directly before the command is called, making
+        it useful for any sort of set up required. This hook is called
+        for both the prefixed and slash versions of the command.
+
+        This pre-invoke hook takes a sole parameter, a :class:`.BridgeContext`.
+
+        Parameters
+        -----------
+        coro: :ref:`coroutine <coroutine>`
+            The coroutine to register as the pre-invoke hook.
+
+        Raises
+        -------
+        TypeError
+            The coroutine passed is not actually a coroutine.
+        """
+
+        if not asyncio.iscoroutinefunction(coro):
+            raise TypeError("The pre-invoke hook must be a coroutine.")
+
+        self.ext_command.before_invoke = coro
+        self.application_command.before_invoke = coro
+
+        return coro
+
+    def after_invoke(self, coro):
+        """A decorator that registers a coroutine as a post-invoke hook.
+
+        This hook is called directly after the command is called, making it
+        useful for any sort of clean up required. This hook is called for
+        both the prefixed and slash versions of the command.
+
+        This post-invoke hook takes a sole parameter, a :class:`.BridgeContext`.
+
+        Parameters
+        -----------
+        coro: :ref:`coroutine <coroutine>`
+            The coroutine to register as the post-invoke hook.
+
+        Raises
+        -------
+        TypeError
+            The coroutine passed is not actually a coroutine.
+        """
+
+        if not asyncio.iscoroutinefunction(coro):
+            raise TypeError("The post-invoke hook must be a coroutine.")
+
+        self.ext_command.after_invoke = coro
+        self.application_command.after_invoke = coro
+
+        return coro
 
 
 def bridge_command(**kwargs):