Enum options (#1292)

Pycord-Development · May 7, 2022 · 7ad078f · 7ad078f
1 parent 6913833
commit 7ad078f
Show file tree

Hide file tree

Showing 2 changed files with 38 additions and 17 deletions.
diff --git a/discord/commands/core.py b/discord/commands/core.py
@@ -32,6 +32,7 @@
 import re
 import types
 from collections import OrderedDict
+from enum import Enum
 from typing import (
     TYPE_CHECKING,
     Any,
@@ -682,25 +683,19 @@ def _parse_options(self, params, *, check_params: bool = True) -> List[Option]:
 
             if self._is_typing_union(option):
                 if self._is_typing_optional(option):
-                    option = Option(option.__args__[0], "No description provided", required=False)  # type: ignore # union type
+                    option = Option(option.__args__[0], required=False)
                 else:
-                    option = Option(option.__args__, "No description provided")  # type: ignore # union type
+                    option = Option(option.__args__)
 
             if not isinstance(option, Option):
-                if isinstance(p_obj.default, Option):  # arg: type = Option(...)
-                    p_obj.default.input_type = SlashCommandOptionType.from_datatype(option)
-                    option = p_obj.default
-                else:  # arg: Option(...) = default
-                    option = Option(option, "No description provided")
-
-            if (
-                option.default is None
-                and p_obj.default != inspect.Parameter.empty
-                and not isinstance(p_obj.default, Option)
-            ):
-                option.default = p_obj.default
-                option.required = False
+                option = Option(option)
 
+            if option.default is None:
+                if p_obj.default == inspect.Parameter.empty:
+                    option.default = None
+                else:
+                    option.default = p_obj.default
+                    option.required = False
             if option.name is None:
                 option.name = p_name
             if option.name != p_name or option._parameter_name is None:
@@ -830,6 +825,15 @@ async def _invoke(self, ctx: ApplicationContext) -> None:
             elif op.input_type == SlashCommandOptionType.string and (converter := op.converter) is not None:
                 arg = await converter.convert(converter, ctx, arg)
 
+            elif issubclass(op._raw_type, Enum):
+                if isinstance(arg, str) and arg.isdigit():
+                    try:
+                        arg = op._raw_type(int(arg))
+                    except ValueError:
+                        arg = op._raw_type(arg)
+                else:
+                    arg = op._raw_type(arg)
+
             kwargs[op._parameter_name] = arg
 
         for o in self.options:

diff --git a/discord/commands/options.py b/discord/commands/options.py
@@ -22,9 +22,11 @@
 DEALINGS IN THE SOFTWARE.
 """
 
+import inspect
 from typing import Any, Dict, List, Literal, Optional, Union
+from enum import Enum
 
-from ..enums import ChannelType, SlashCommandOptionType
+from ..enums import ChannelType, SlashCommandOptionType, Enum as DiscordEnum
 
 __all__ = (
     "ThreadOption",
@@ -124,10 +126,18 @@ def __init__(self, input_type: Any = str, /, description: Optional[str] = None,
         self.converter = None
         self._raw_type = input_type
         self.channel_types: List[ChannelType] = kwargs.pop("channel_types", [])
+        enum_choices = []
         if not isinstance(input_type, SlashCommandOptionType):
             if hasattr(input_type, "convert"):
                 self.converter = input_type
                 input_type = SlashCommandOptionType.string
+            elif issubclass(input_type, (Enum, DiscordEnum)):
+                enum_choices = [OptionChoice(e.name, e.value) for e in input_type]
+                if len(enum_choices) != len([elem for elem in enum_choices if elem.value.__class__ == enum_choices[0].value.__class__]):
+                    enum_choices = [OptionChoice(e.name, str(e.value)) for e in input_type]
+                    input_type = SlashCommandOptionType.string
+                else:
+                    input_type = SlashCommandOptionType.from_datatype(enum_choices[0].value.__class__)
             else:
                 try:
                     _type = SlashCommandOptionType.from_datatype(input_type)
@@ -158,10 +168,17 @@ def __init__(self, input_type: Any = str, /, description: Optional[str] = None,
         self.input_type = input_type
         self.required: bool = kwargs.pop("required", True) if "default" not in kwargs else False
         self.default = kwargs.pop("default", None)
-        self.choices: List[OptionChoice] = [
+        self.choices: List[OptionChoice] = enum_choices or [
             o if isinstance(o, OptionChoice) else OptionChoice(o) for o in kwargs.pop("choices", list())
         ]
 
+        if description is not None:
+            self.description = description
+        elif issubclass(self._raw_type, Enum) and (doc := inspect.getdoc(self._raw_type)) is not None:
+            self.description = doc
+        else:
+            self.description = "No description provided"
+
         if self.input_type == SlashCommandOptionType.integer:
             minmax_types = (int, type(None))
         elif self.input_type == SlashCommandOptionType.number: