From 7ad078fc778b1b86210cf7879b12dd7ce448e188 Mon Sep 17 00:00:00 2001 From: Middledot <78228142+Middledot@users.noreply.github.com> Date: Sat, 7 May 2022 17:08:47 -0400 Subject: [PATCH] Enum options (#1292) --- discord/commands/core.py | 34 +++++++++++++++++++--------------- discord/commands/options.py | 21 +++++++++++++++++++-- 2 files changed, 38 insertions(+), 17 deletions(-) diff --git a/discord/commands/core.py b/discord/commands/core.py index 4ec6982321..7e5aa728e3 100644 --- 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 index e60c53c440..13156325fc 100644 --- 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: