Add support for autocomplete

This commit is contained in:
Rapptz
2022-03-01 04:38:41 -05:00
parent 4e04dbdec7
commit ae1aaac5a7
10 changed files with 273 additions and 38 deletions

View File

@ -14,5 +14,5 @@ from .enums import *
from .errors import *
from .models import *
from .tree import *
from .namespace import Namespace
from .namespace import *
from .transformers import *

View File

@ -37,39 +37,27 @@ from typing import (
Set,
TYPE_CHECKING,
Tuple,
Type,
TypeVar,
Union,
)
from textwrap import TextWrapper
import sys
import re
from .enums import AppCommandOptionType, AppCommandType
from ..interactions import Interaction
from ..enums import ChannelType, try_enum
from .models import AppCommandChannel, AppCommandThread, Choice
from .models import Choice
from .transformers import annotation_to_parameter, CommandParameter, NoneType
from .errors import AppCommandError, CommandInvokeError, CommandSignatureMismatch, CommandAlreadyRegistered
from ..utils import resolve_annotation, MISSING, is_inside_class
from ..user import User
from ..member import Member
from ..role import Role
from ..message import Message
from ..mixins import Hashable
from ..permissions import Permissions
if TYPE_CHECKING:
from typing_extensions import ParamSpec, Concatenate
from ..types.interactions import (
ResolvedData,
PartialThread,
PartialChannel,
ApplicationCommandInteractionDataOption,
)
from ..state import ConnectionState
from ..user import User
from ..member import Member
from ..message import Message
from .namespace import Namespace
from .models import ChoiceT
__all__ = (
'Command',
@ -93,25 +81,33 @@ Error = Union[
Callable[[Interaction, AppCommandError], Coro[Any]],
]
ContextMenuCallback = Union[
# If groups end up support context menus these would be uncommented
# Callable[[GroupT, Interaction, Member], Coro[Any]],
# Callable[[GroupT, Interaction, User], Coro[Any]],
# Callable[[GroupT, Interaction, Message], Coro[Any]],
# Callable[[GroupT, Interaction, Union[Member, User]], Coro[Any]],
Callable[[Interaction, Member], Coro[Any]],
Callable[[Interaction, User], Coro[Any]],
Callable[[Interaction, Message], Coro[Any]],
Callable[[Interaction, Union[Member, User]], Coro[Any]],
]
if TYPE_CHECKING:
CommandCallback = Union[
Callable[Concatenate[GroupT, Interaction, P], Coro[T]],
Callable[Concatenate[Interaction, P], Coro[T]],
]
ContextMenuCallback = Union[
# If groups end up support context menus these would be uncommented
# Callable[[GroupT, Interaction, Member], Coro[Any]],
# Callable[[GroupT, Interaction, User], Coro[Any]],
# Callable[[GroupT, Interaction, Message], Coro[Any]],
# Callable[[GroupT, Interaction, Union[Member, User]], Coro[Any]],
Callable[[Interaction, Member], Coro[Any]],
Callable[[Interaction, User], Coro[Any]],
Callable[[Interaction, Message], Coro[Any]],
Callable[[Interaction, Union[Member, User]], Coro[Any]],
]
AutocompleteCallback = Union[
Callable[[GroupT, Interaction, ChoiceT, Namespace], Coro[List[Choice[ChoiceT]]]],
Callable[[Interaction, ChoiceT, Namespace], Coro[List[Choice[ChoiceT]]]],
]
else:
CommandCallback = Callable[..., Coro[T]]
ContextMenuCallback = Callable[..., Coro[T]]
AutocompleteCallback = Callable[..., Coro[T]]
VALID_SLASH_COMMAND_NAME = re.compile(r'^[\w-]{1,32}$')
@ -197,6 +193,25 @@ def _populate_choices(params: Dict[str, CommandParameter], all_choices: Dict[str
raise TypeError(f'unknown parameter given: {first}')
def _populate_autocomplete(params: Dict[str, CommandParameter], autocomplete: Dict[str, Any]) -> None:
for name, param in params.items():
callback = autocomplete.pop(name, MISSING)
if callback is MISSING:
continue
if not inspect.iscoroutinefunction(callback):
raise TypeError('autocomplete callback must be a coroutine function')
if param.type not in (AppCommandOptionType.string, AppCommandOptionType.number, AppCommandOptionType.integer):
raise TypeError('autocomplete is only supported for integer, string, or number option types')
param.autocomplete = callback
if autocomplete:
first = next(iter(autocomplete))
raise TypeError(f'unknown parameter given: {first}')
def _extract_parameters_from_callback(func: Callable[..., Any], globalns: Dict[str, Any]) -> Dict[str, CommandParameter]:
params = inspect.signature(func).parameters
cache = {}
@ -236,6 +251,13 @@ def _extract_parameters_from_callback(func: Callable[..., Any], globalns: Dict[s
else:
_populate_choices(result, choices)
try:
autocomplete = func.__discord_app_commands_param_autocomplete__
except AttributeError:
pass
else:
_populate_autocomplete(result, autocomplete)
return result
@ -381,6 +403,27 @@ class Command(Generic[GroupT, P, T]):
except Exception as e:
raise CommandInvokeError(self, e) from e
async def _invoke_autocomplete(self, interaction: Interaction, name: str, namespace: Namespace):
value = namespace.__dict__[name]
try:
param = self._params[name]
except KeyError:
raise CommandSignatureMismatch(self) from None
if param.autocomplete is None:
raise CommandSignatureMismatch(self)
if self.binding is not None:
choices = await param.autocomplete(self.binding, interaction, value, namespace)
else:
choices = await param.autocomplete(interaction, value, namespace)
if interaction.response.is_done():
return
await interaction.response.autocomplete(choices)
def _get_internal_command(self, name: str) -> Optional[Union[Command, Group]]:
return None
@ -418,6 +461,69 @@ class Command(Generic[GroupT, P, T]):
self.on_error = coro
return coro
def autocomplete(
self, name: str
) -> Callable[[AutocompleteCallback[GroupT, ChoiceT]], AutocompleteCallback[GroupT, ChoiceT]]:
"""A decorator that registers a coroutine as an autocomplete prompt for a parameter.
The coroutine callback must have 3 parameters, the :class:`~discord.Interaction`,
the current value by the user (usually either a :class:`str`, :class:`int`, or :class:`float`,
depending on the type of the parameter being marked as autocomplete), and then the
:class:`Namespace` that represents possible values are partially filled in.
The coroutine decorator **must** return a list of :class:`~discord.app_commands.Choice` objects.
Only up to 25 objects are supported.
Example:
.. code-block:: python3
@app_commands.command()
async def fruits(interaction: discord.Interaction, fruits: str):
await interaction.response.send_message(f'Your favourite fruit seems to be {fruits}')
@fruits.autocomplete('fruits')
async def fruits_autocomplete(
interaction: discord.Interaction,
current: str,
namespace: app_commands.Namespace
) -> List[app_commands.Choice[str]]:
fruits = ['Banana', 'Pineapple', 'Apple', 'Watermelon', 'Melon', 'Cherry']
return [
app_commands.Choice(name=fruit, value=fruit)
for fruit in fruits if current.lower() in fruit.lower()
]
Parameters
-----------
name: :clas:`str`
The parameter name to register as autocomplete.
Raises
-------
TypeError
The coroutine passed is not actually a coroutine or
the parameter is not found or of an invalid type.
"""
def decorator(coro: AutocompleteCallback[GroupT, ChoiceT]) -> AutocompleteCallback[GroupT, ChoiceT]:
if not inspect.iscoroutinefunction(coro):
raise TypeError('The error handler must be a coroutine.')
try:
param = self._params[name]
except KeyError:
raise TypeError(f'unknown parameter: {name!r}') from None
if param.type not in (AppCommandOptionType.string, AppCommandOptionType.number, AppCommandOptionType.integer):
raise TypeError('autocomplete is only supported for integer, string, or number option types')
param.autocomplete = coro
return coro
return decorator
class ContextMenu:
"""A class that implements a context menu application command.
@ -882,7 +988,7 @@ def choices(**parameters: List[Choice]) -> Callable[[T], T]:
Raises
--------
TypeError
The parameter name is not found.
The parameter name is not found or the parameter type was incorrect.
"""
def decorator(inner: T) -> T:
@ -897,3 +1003,54 @@ def choices(**parameters: List[Choice]) -> Callable[[T], T]:
return inner
return decorator
def autocomplete(**parameters: AutocompleteCallback[GroupT, ChoiceT]) -> Callable[[T], T]:
r"""Associates the given parameters with the given autocomplete callback.
Autocomplete is only supported on types that have :class:`str`, :class:`int`, or :class:`float`
values.
Example:
.. code-block:: python3
@app_commands.command()
@app_commands.autocomplete(fruits=fruits_autocomplete)
async def fruits(interaction: discord.Interaction, fruits: str):
await interaction.response.send_message(f'Your favourite fruit seems to be {fruits}')
async def fruits_autocomplete(
interaction: discord.Interaction,
current: str,
namespace: app_commands.Namespace
) -> List[app_commands.Choice[str]]:
fruits = ['Banana', 'Pineapple', 'Apple', 'Watermelon', 'Melon', 'Cherry']
return [
app_commands.Choice(name=fruit, value=fruit)
for fruit in fruits if current.lower() in fruit.lower()
]
Parameters
-----------
\*\*parameters
The parameters to mark as autocomplete.
Raises
--------
TypeError
The parameter name is not found or the parameter type was incorrect.
"""
def decorator(inner: T) -> T:
if isinstance(inner, Command):
_populate_autocomplete(inner._params, parameters)
else:
try:
inner.__discord_app_commands_param_autocomplete__.update(parameters) # type: ignore - Runtime attribute access
except AttributeError:
inner.__discord_app_commands_param_autocomplete__ = parameters # type: ignore - Runtime attribute assignment
return inner
return decorator

View File

@ -31,7 +31,7 @@ from ..enums import ChannelType, try_enum
from ..mixins import Hashable
from ..utils import _get_as_snowflake, parse_time, snowflake_time
from .enums import AppCommandOptionType, AppCommandType
from typing import Generic, List, NamedTuple, TYPE_CHECKING, Optional, TypeVar, Union
from typing import Generic, List, TYPE_CHECKING, Optional, TypeVar, Union
__all__ = (
'AppCommand',

View File

@ -37,6 +37,8 @@ from .enums import AppCommandOptionType
if TYPE_CHECKING:
from ..types.interactions import ResolvedData, ApplicationCommandInteractionDataOption
__all__ = ('Namespace',)
class ResolveKey(NamedTuple):
id: str

View File

@ -27,7 +27,22 @@ import inspect
from dataclasses import dataclass
from enum import Enum
from typing import TYPE_CHECKING, Any, ClassVar, Dict, List, Literal, Optional, Set, Tuple, Type, TypeVar, Union
from typing import (
TYPE_CHECKING,
Any,
Callable,
ClassVar,
Coroutine,
Dict,
List,
Literal,
Optional,
Set,
Tuple,
Type,
TypeVar,
Union,
)
from .enums import AppCommandOptionType
from .errors import TransformerError
@ -75,8 +90,6 @@ class CommandParameter:
The minimum supported value for this parameter.
max_value: Optional[Union[:class:`int`, :class:`float`]]
The maximum supported value for this parameter.
autocomplete: :class:`bool`
Whether this parameter enables autocomplete.
"""
name: str = MISSING
@ -88,7 +101,7 @@ class CommandParameter:
channel_types: List[ChannelType] = MISSING
min_value: Optional[Union[int, float]] = None
max_value: Optional[Union[int, float]] = None
autocomplete: bool = MISSING
autocomplete: Optional[Callable[..., Coroutine[Any, Any, Any]]] = None
_annotation: Any = MISSING
def to_dict(self) -> Dict[str, Any]:

View File

@ -26,7 +26,7 @@ from __future__ import annotations
import inspect
import sys
import traceback
from typing import Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Tuple, Type, Union, overload
from typing import Callable, Dict, List, Literal, Optional, TYPE_CHECKING, Tuple, Union, overload
from .namespace import Namespace, ResolveKey
@ -40,6 +40,7 @@ from .errors import (
CommandSignatureMismatch,
)
from ..errors import ClientException
from ..enums import InteractionType
from ..utils import MISSING
if TYPE_CHECKING:
@ -580,7 +581,7 @@ class CommandTree:
raise CommandSignatureMismatch(ctx_menu)
if value is None:
raise RuntimeError('This should not happen if Discord sent well-formed data.')
raise AppCommandError('This should not happen if Discord sent well-formed data.')
# I assume I don't have to type check here.
try:
@ -608,6 +609,8 @@ class CommandTree:
CommandSignatureMismatch
The interaction data referred to a parameter that was not found in the
application command definition.
AppCommandError
An error occurred while calling the command.
"""
data: ApplicationCommandInteractionData = interaction.data # type: ignore
type = data.get('type', 1)
@ -663,6 +666,14 @@ class CommandTree:
# and command refers to the class type we care about
namespace = Namespace(interaction, data.get('resolved', {}), options)
# Auto complete handles the namespace differently... so at this point this is where we decide where that is.
if interaction.type is InteractionType.autocomplete:
focused = next((opt['name'] for opt in options if opt.get('focused')), None)
if focused is None:
raise AppCommandError('This should not happen, but there is no focused element. This is a Discord bug.')
await command._invoke_autocomplete(interaction, focused, namespace)
return
try:
await command._invoke_with_namespace(interaction, namespace)
except AppCommandError as e: