Merge pull request #285 from goverfl0w/master

Implement context menus.

Author: i0bs

discord_slash/__init__.py

@@ -1,17 +1,15 @@
 """
 discord-py-slash-command
 ~~~~~~~~~~~~~~~~~~~~~~~~
-
 Simple Discord Slash Command extension for discord.py
-
 :copyright: (c) 2020-2021 eunwoo1104
 :license: MIT
 """
 
 from .client import SlashCommand  # noqa: F401
 from .const import __version__  # noqa: F401
-from .context import ComponentContext, SlashContext  # noqa: F401
+from .context import ComponentContext, MenuContext, SlashContext  # noqa: F401
 from .dpy_overrides import ComponentMessage  # noqa: F401
-from .model import ButtonStyle, ComponentType, SlashCommandOptionType  # noqa: F401
+from .model import ButtonStyle, ComponentType, ContextMenuType, SlashCommandOptionType  # noqa: F401
 from .utils import manage_commands  # noqa: F401
 from .utils import manage_components  # noqa: F401

discord_slash/client.py

@@ -46,6 +46,7 @@ class SlashCommand:
 
     :ivar _discord: Discord client of this client.
     :ivar commands: Dictionary of the registered commands via :func:`.slash` decorator.
+    :ivar menu_commands: Dictionary of the registered context menus via the :func:`.context_menu` decorator.
     :ivar req: :class:`.http.SlashCommandRequest` of this client.
     :ivar logger: Logger of this client.
     :ivar sync_commands: Whether to sync commands automatically.
@@ -64,7 +65,7 @@ def __init__(
         application_id: typing.Optional[int] = None,
     ):
         self._discord = client
-        self.commands = {}
+        self.commands = {"context": {}}
         self.subcommands = {}
         self.components = {}
         self.logger = logging.getLogger("discord_slash")
@@ -270,12 +271,53 @@ async def to_dict(self):
         await self._discord.wait_until_ready()  # In case commands are still not registered to SlashCommand.
         all_guild_ids = []
         for x in self.commands:
+            if x == "context":
+                # handle context menu separately.
+                for _x in self.commands["context"]:
+                    _selected = self.commands["context"][_x]
+                    for i in _selected.allowed_guild_ids:
+                        if i not in all_guild_ids:
+                            all_guild_ids.append(i)
+                continue
             for i in self.commands[x].allowed_guild_ids:
                 if i not in all_guild_ids:
                     all_guild_ids.append(i)
         cmds = {"global": [], "guild": {x: [] for x in all_guild_ids}}
         wait = {}  # Before merging to return dict, let's first put commands to temporary dict.
         for x in self.commands:
+            if x == "context":
+                # handle context menu separately.
+                for _x in self.commands["context"]:  # x is the new reference dict
+                    selected = self.commands["context"][_x]
+
+                    if selected.allowed_guild_ids:
+                        for y in selected.allowed_guild_ids:
+                            if y not in wait:
+                                wait[y] = {}
+                            command_dict = {
+                                "name": _x,
+                                "options": selected.options or [],
+                                "default_permission": selected.default_permission,
+                                "permissions": {},
+                                "type": selected._type,
+                            }
+                            if y in selected.permissions:
+                                command_dict["permissions"][y] = selected.permissions[y]
+                            wait[y][x] = copy.deepcopy(command_dict)
+                    else:
+                        if "global" not in wait:
+                            wait["global"] = {}
+                        command_dict = {
+                            "name": _x,
+                            "options": selected.options or [],
+                            "default_permission": selected.default_permission,
+                            "permissions": selected.permissions or {},
+                            "type": selected._type,
+                        }
+                        wait["global"][x] = copy.deepcopy(command_dict)
+
+                continue
+
             selected = self.commands[x]
             if selected.allowed_guild_ids:
                 for y in selected.allowed_guild_ids:
@@ -287,7 +329,10 @@ async def to_dict(self):
                         "options": selected.options or [],
                         "default_permission": selected.default_permission,
                         "permissions": {},
+                        "type": selected._type,
                     }
+                    if command_dict["type"] != 1:
+                        command_dict.pop("description")
                     if y in selected.permissions:
                         command_dict["permissions"][y] = selected.permissions[y]
                     wait[y][x] = copy.deepcopy(command_dict)
@@ -300,14 +345,20 @@ async def to_dict(self):
                     "options": selected.options or [],
                     "default_permission": selected.default_permission,
                     "permissions": selected.permissions or {},
+                    "type": selected._type,
                 }
+                if command_dict["type"] != 1:
+                    command_dict.pop("description")
                 wait["global"][x] = copy.deepcopy(command_dict)
 
         # Separated normal command add and subcommand add not to
         # merge subcommands to one. More info at Issue #88
         # https://github.com/eunwoo1104/discord-py-slash-command/issues/88
 
         for x in self.commands:
+            if x == "context":
+                continue  # no menus have subcommands.
+
             if not self.commands[x].has_subcommands:
                 continue
             tgt = self.subcommands[x]
@@ -424,7 +475,7 @@ async def sync_all_commands(
                     if ex.status == 400:
                         # catch bad requests
                         cmd_nums = set(
-                            re.findall(r"In\s(\d).", ex.args[0])
+                            re.findall(r"^[\w-]{1,32}$", ex.args[0])
                         )  # find all discords references to commands
                         error_string = ex.args[0]
 
@@ -594,6 +645,66 @@ def add_slash_command(
         self.logger.debug(f"Added command `{name}`")
         return obj
 
+    def _cog_ext_add_context_menu(self, target: int, name: str, guild_ids: list = None):
+        """
+        Creates a new cog_based context menu command.
+
+        :param cmd: Command Coroutine.
+        :type cmd: Coroutine
+        :param name: The name of the command
+        :type name: str
+        :param _type: The context menu type.
+        :type _type: int
+        """
+
+    def add_context_menu(self, cmd, name: str, _type: int, guild_ids: list = None):
+        """
+        Creates a new context menu command.
+
+        :param cmd: Command Coroutine.
+        :type cmd: Coroutine
+        :param name: The name of the command
+        :type name: str
+        :param _type: The context menu type.
+        :type _type: int
+        """
+
+        name = [name or cmd.__name__][0]
+        guild_ids = guild_ids or []
+
+        if not all(isinstance(item, int) for item in guild_ids):
+            raise error.IncorrectGuildIDType(
+                f"The snowflake IDs {guild_ids} given are not a list of integers. Because of discord.py convention, please use integer IDs instead. Furthermore, the command '{name}' will be deactivated and broken until fixed."
+            )
+
+        if name in self.commands["context"]:
+            tgt = self.commands["context"][name]
+            if not tgt.has_subcommands:
+                raise error.DuplicateCommand(name)
+            has_subcommands = tgt.has_subcommands  # noqa
+            for x in tgt.allowed_guild_ids:
+                if x not in guild_ids:
+                    guild_ids.append(x)
+
+        _cmd = {
+            "default_permission": None,
+            "has_permissions": None,
+            "name": name,
+            "type": _type,
+            "func": cmd,
+            "description": "",
+            "guild_ids": guild_ids,
+            "api_options": [],
+            "connector": {},
+            "has_subcommands": False,
+            "api_permissions": {},
+        }
+
+        obj = model.BaseCommandObject(name, cmd=_cmd, _type=_type)
+        self.commands["context"][name] = obj
+        self.logger.debug(f"Added context command `{name}`")
+        return obj
+
     def add_subcommand(
         self,
         cmd,
@@ -916,6 +1027,34 @@ def wrapper(cmd):
 
         return wrapper
 
+    def context_menu(self, *, target: int, name: str, guild_ids: list = None):
+        """
+        Decorator that adds context menu commands.
+
+        :param target: The type of menu.
+        :type target: int
+        :param name: A name to register as the command in the menu.
+        :type name: str
+        :param guild_ids: A list of guild IDs to register the command under. Defaults to ``None``.
+        :type guild_ids: list
+        """
+
+        def wrapper(cmd):
+            # _obj = self.add_slash_command(
+            #    cmd,
+            #    name,
+            #    "",
+            #    guild_ids
+            # )
+
+            # This has to call both, as its a arg-less menu.
+
+            obj = self.add_context_menu(cmd, name, target, guild_ids)
+
+            return obj
+
+        return wrapper
+
     def add_component_callback(
         self,
         callback: typing.Coroutine,
@@ -1255,12 +1394,15 @@ async def on_socket_response(self, msg):
 
         to_use = msg["d"]
         interaction_type = to_use["type"]
-        if interaction_type in (1, 2):
-            return await self._on_slash(to_use)
-        if interaction_type == 3:
-            return await self._on_component(to_use)
-
-        raise NotImplementedError
+        if interaction_type in (1, 2, 3) or msg["s"] == 5:
+            await self._on_slash(to_use)
+            await self._on_context_menu(to_use)
+            try:
+                await self._on_component(to_use)  # noqa
+            except KeyError:
+                pass  # for some reason it complains about custom_id being an optional arg when it's fine?
+        return
+        # raise NotImplementedError
 
     async def _on_component(self, to_use):
         ctx = context.ComponentContext(self.req, to_use, self._discord, self.logger)
@@ -1319,6 +1461,34 @@ async def _on_slash(self, to_use):
 
             await self.invoke_command(selected_cmd, ctx, args)
 
+    async def _on_context_menu(self, to_use):
+        if to_use["data"]["name"] in self.commands["context"]:
+            ctx = context.MenuContext(self.req, to_use, self._discord, self.logger)
+            cmd_name = to_use["data"]["name"]
+
+            if cmd_name not in self.commands["context"] and cmd_name in self.subcommands:
+                return  # menus don't have subcommands you smooth brain
+
+            selected_cmd = self.commands["context"][cmd_name]
+
+            if (
+                selected_cmd.allowed_guild_ids
+                and ctx.guild_id not in selected_cmd.allowed_guild_ids
+            ):
+                return
+
+            if selected_cmd.has_subcommands and not selected_cmd.func:
+                return await self.handle_subcommand(ctx, to_use)
+
+            if "options" in to_use["data"]:
+                for x in to_use["data"]["options"]:
+                    if "value" not in x:
+                        return await self.handle_subcommand(ctx, to_use)
+
+            self._discord.dispatch("context_menu", ctx)
+
+            await self.invoke_command(selected_cmd, ctx, args={})
+
     async def handle_subcommand(self, ctx: context.SlashContext, data: dict):
         """
         Coroutine for handling subcommand.

discord_slash/cog_ext.py

@@ -189,6 +189,49 @@ def wrapper(cmd):
     return wrapper
 
 
+# I don't feel comfortable with having these right now, they're too buggy even when they were working.
+
+
+def cog_context_menu(*, name: str, guild_ids: list = None, target: int = 1):
+    """
+    Decorator that adds context menu commands.
+
+    :param target: The type of menu.
+    :type target: int
+    :param name: A name to register as the command in the menu.
+    :type name: str
+    :param guild_ids: A list of guild IDs to register the command under. Defaults to ``None``.
+    :type guild_ids: list
+    """
+
+    def wrapper(cmd):
+        # _obj = self.add_slash_command(
+        #    cmd,
+        #    name,
+        #    "",
+        #    guild_ids
+        # )
+
+        # This has to call both, as its a arg-less menu.
+
+        _cmd = {
+            "default_permission": None,
+            "has_permissions": None,
+            "name": name,
+            "type": target,
+            "func": cmd,
+            "description": "",
+            "guild_ids": guild_ids,
+            "api_options": [],
+            "connector": {},
+            "has_subcommands": False,
+            "api_permissions": {},
+        }
+        return CogBaseCommandObject(name or cmd.__name__, _cmd, target)
+
+    return wrapper
+
+
 def permission(guild_id: int, permissions: list):
     """
     Decorator that add permissions. This will set the permissions for a single guild, you can use it more than once for each command.