From 375a001da8437c4d6359b925ea9b3502d802839a Mon Sep 17 00:00:00 2001 From: James Hilton-Balfe Date: Tue, 17 Jan 2023 01:36:00 +0000 Subject: [PATCH] Add overloads for all the events in the event reference Co-authored-by: Sebastian Law Co-authored-by: Josh <8677174+bijij@users.noreply.github.com> --- discord/client.py | 739 +++++++++++++++++++++++++++++++++++++++++++++- 1 file changed, 730 insertions(+), 9 deletions(-) diff --git a/discord/client.py b/discord/client.py index 3878ccd0d..b7333ff02 100644 --- a/discord/client.py +++ b/discord/client.py @@ -28,6 +28,7 @@ import asyncio import datetime import logging from typing import ( + TYPE_CHECKING, Any, AsyncIterator, Callable, @@ -35,13 +36,14 @@ from typing import ( Dict, Generator, List, + Literal, Optional, Sequence, - TYPE_CHECKING, Tuple, Type, TypeVar, Union, + overload, ) import aiohttp @@ -75,23 +77,52 @@ from .threads import Thread from .sticker import GuildSticker, StandardSticker, StickerPack, _sticker_factory if TYPE_CHECKING: - from typing_extensions import Self from types import TracebackType - from .types.guild import Guild as GuildPayload - from .abc import SnowflakeTime, Snowflake, PrivateChannel + + from typing_extensions import Self + + from .abc import Messageable, PrivateChannel, Snowflake, SnowflakeTime + from .app_commands import Command, ContextMenu + from .automod import AutoModAction, AutoModRule + from .channel import DMChannel, GroupChannel + from .ext.commands import AutoShardedBot, Bot, Context, CommandError from .guild import GuildChannel - from .channel import DMChannel + from .integrations import Integration + from .interactions import Interaction + from .member import Member, VoiceState from .message import Message - from .member import Member + from .raw_models import ( + RawAppCommandPermissionsUpdateEvent, + RawBulkMessageDeleteEvent, + RawIntegrationDeleteEvent, + RawMemberRemoveEvent, + RawMessageDeleteEvent, + RawMessageUpdateEvent, + RawReactionActionEvent, + RawReactionClearEmojiEvent, + RawReactionClearEvent, + RawThreadDeleteEvent, + RawThreadMembersUpdate, + RawThreadUpdateEvent, + RawTypingEvent, + ) + from .reaction import Reaction + from .role import Role + from .scheduled_event import ScheduledEvent + from .threads import ThreadMember + from .types.guild import Guild as GuildPayload from .voice_client import VoiceProtocol + # fmt: off __all__ = ( 'Client', ) # fmt: on -Coro = TypeVar('Coro', bound=Callable[..., Coroutine[Any, Any, Any]]) +T = TypeVar('T') +Coro = Coroutine[Any, Any, T] +CoroT = TypeVar('CoroT', bound=Callable[..., Coro[Any]]) _log = logging.getLogger(__name__) @@ -1110,7 +1141,687 @@ class Client: 'Please use the login method or asynchronous context manager before calling this method' ) - def wait_for( + # App Commands + + @overload + async def wait_for( + self, + event: Literal['raw_app_command_permissions_update'], + /, + *, + check: Optional[Callable[[RawAppCommandPermissionsUpdateEvent], bool]], + timeout: Optional[float] = None, + ) -> RawAppCommandPermissionsUpdateEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['app_command_completion'], + /, + *, + check: Optional[Callable[[Interaction, Union[Command, ContextMenu]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Interaction, Union[Command, ContextMenu]]: + ... + + # AutoMod + + @overload + async def wait_for( + self, + event: Literal['automod_rule_create', 'automod_rule_update', 'automod_rule_delete'], + /, + *, + check: Optional[Callable[[AutoModRule], bool]], + timeout: Optional[float] = None, + ) -> AutoModRule: + ... + + @overload + async def wait_for( + self, + event: Literal['automod_action'], + /, + *, + check: Optional[Callable[[AutoModAction], bool]], + timeout: Optional[float] = None, + ) -> AutoModAction: + ... + + # Channels + + @overload + async def wait_for( + self, + event: Literal['private_channel_update'], + /, + *, + check: Optional[Callable[[GroupChannel, GroupChannel], bool]], + timeout: Optional[float] = None, + ) -> Tuple[GroupChannel, GroupChannel]: + ... + + @overload + async def wait_for( + self, + event: Literal['private_channel_pins_update'], + /, + *, + check: Optional[Callable[[PrivateChannel, datetime.datetime], bool]], + timeout: Optional[float] = None, + ) -> Tuple[PrivateChannel, datetime.datetime]: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_channel_delete', 'guild_channel_create'], + /, + *, + check: Optional[Callable[[GuildChannel], bool]], + timeout: Optional[float] = None, + ) -> GuildChannel: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_channel_update'], + /, + *, + check: Optional[Callable[[GuildChannel, GuildChannel], bool]], + timeout: Optional[float] = None, + ) -> Tuple[GuildChannel, GuildChannel]: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_channel_pins_update'], + /, + *, + check: Optional[ + Callable[ + [Union[GuildChannel, Thread], Optional[datetime.datetime]], + bool, + ] + ], + timeout: Optional[float] = None, + ) -> Tuple[Union[GuildChannel, Thread], Optional[datetime.datetime]]: + ... + + @overload + async def wait_for( + self, + event: Literal['typing'], + /, + *, + check: Optional[Callable[[Messageable, Union[User, Member], datetime.datetime], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Messageable, Union[User, Member], datetime.datetime]: + ... + + @overload + async def wait_for( + self, + event: Literal['on_raw_typing'], + /, + *, + check: Optional[Callable[[RawTypingEvent], bool]], + timeout: Optional[float] = None, + ) -> RawTypingEvent: + ... + + # Debug & Gateway events + + @overload + async def wait_for( + self, + event: Literal['connect', 'disconnect', 'ready', 'resumed'], + /, + *, + check: Optional[Callable[[], bool]], + timeout: Optional[float] = None, + ) -> None: + ... + + @overload + async def wait_for( + self, + event: Literal['shard_connect', 'shard_disconnect', 'shard_ready', 'shard_resumed'], + /, + *, + check: Optional[Callable[[int], bool]], + timeout: Optional[float] = None, + ) -> int: + ... + + @overload + async def wait_for( + self, + event: Literal['socket_event_type', 'socket_raw_receive'], + /, + *, + check: Optional[Callable[[str], bool]], + timeout: Optional[float] = None, + ) -> str: + ... + + @overload + async def wait_for( + self, + event: Literal['socket_raw_send'], + /, + *, + check: Optional[Callable[[Union[str, bytes]], bool]], + timeout: Optional[float] = None, + ) -> Union[str, bytes]: + ... + + # Guilds + + @overload + async def wait_for( + self, + event: Literal[ + 'guild_available', + 'guild_unavailable', + 'guild_join', + 'guild_remove', + ], + /, + *, + check: Optional[Callable[[Guild], bool]], + timeout: Optional[float] = None, + ) -> Guild: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_update'], + /, + *, + check: Optional[Callable[[Guild, Guild], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Guild, Guild]: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_emojis_update'], + /, + *, + check: Optional[Callable[[Guild, Sequence[Emoji], Sequence[Emoji]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Guild, Sequence[Emoji], Sequence[Emoji]]: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_stickers_update'], + /, + *, + check: Optional[Callable[[Guild, Sequence[GuildSticker], Sequence[GuildSticker]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Guild, Sequence[GuildSticker], Sequence[GuildSticker]]: + ... + + @overload + async def wait_for( + self, + event: Literal['invite_create', 'invite_delete'], + /, + *, + check: Optional[Callable[[Invite], bool]], + timeout: Optional[float] = None, + ) -> Invite: + ... + + # Integrations + + @overload + async def wait_for( + self, + event: Literal['integration_create', 'integration_update'], + /, + *, + check: Optional[Callable[[Integration], bool]], + timeout: Optional[float] = None, + ) -> Integration: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_integrations_update'], + /, + *, + check: Optional[Callable[[Guild], bool]], + timeout: Optional[float] = None, + ) -> Guild: + ... + + @overload + async def wait_for( + self, + event: Literal['webhooks_update'], + /, + *, + check: Optional[Callable[[GuildChannel], bool]], + timeout: Optional[float] = None, + ) -> GuildChannel: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_integration_delete'], + /, + *, + check: Optional[Callable[[RawIntegrationDeleteEvent], bool]], + timeout: Optional[float] = None, + ) -> RawIntegrationDeleteEvent: + ... + + # Interactions + + @overload + async def wait_for( + self, + event: Literal['interaction'], + /, + *, + check: Optional[Callable[[Interaction], bool]], + timeout: Optional[float] = None, + ) -> Interaction: + ... + + # Members + + @overload + async def wait_for( + self, + event: Literal['member_join', 'member_remove'], + /, + *, + check: Optional[Callable[[Member], bool]], + timeout: Optional[float] = None, + ) -> Member: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_member_remove'], + /, + *, + check: Optional[Callable[[RawMemberRemoveEvent], bool]], + timeout: Optional[float] = None, + ) -> RawMemberRemoveEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['member_update', 'presence_update'], + /, + *, + check: Optional[Callable[[Member, Member], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Member, Member]: + ... + + @overload + async def wait_for( + self, + event: Literal['user_update'], + /, + *, + check: Optional[Callable[[User, User], bool]], + timeout: Optional[float] = None, + ) -> Tuple[User, User]: + ... + + @overload + async def wait_for( + self, + event: Literal['member_ban'], + /, + *, + check: Optional[Callable[[Guild, Union[User, Member]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Guild, Union[User, Member]]: + ... + + @overload + async def wait_for( + self, + event: Literal['member_unban'], + /, + *, + check: Optional[Callable[[Guild, User], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Guild, User]: + ... + + # Messages + + @overload + async def wait_for( + self, + event: Literal['message', 'message_delete'], + /, + *, + check: Optional[Callable[[Message], bool]], + timeout: Optional[float] = None, + ) -> Message: + ... + + @overload + async def wait_for( + self, + event: Literal['message_edit'], + /, + *, + check: Optional[Callable[[Message, Message], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Message, Message]: + ... + + @overload + async def wait_for( + self, + event: Literal['bulk_message_delete'], + /, + *, + check: Optional[Callable[[List[Message]], bool]], + timeout: Optional[float] = None, + ) -> List[Message]: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_message_edit'], + /, + *, + check: Optional[Callable[[RawMessageUpdateEvent], bool]], + timeout: Optional[float] = None, + ) -> RawMessageUpdateEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_message_delete'], + /, + *, + check: Optional[Callable[[RawMessageDeleteEvent], bool]], + timeout: Optional[float] = None, + ) -> RawMessageDeleteEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_bulk_message_delete'], + /, + *, + check: Optional[Callable[[RawBulkMessageDeleteEvent], bool]], + timeout: Optional[float] = None, + ) -> RawBulkMessageDeleteEvent: + ... + + # Reactions + + @overload + async def wait_for( + self, + event: Literal['reaction_add', 'reaction_remove'], + /, + *, + check: Optional[Callable[[Reaction, Union[Member, User]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Reaction, Union[Member, User]]: + ... + + @overload + async def wait_for( + self, + event: Literal['reaction_clear'], + /, + *, + check: Optional[Callable[[Message, List[Reaction]], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Message, List[Reaction]]: + ... + + @overload + async def wait_for( + self, + event: Literal['reaction_clear_emoji'], + /, + *, + check: Optional[Callable[[Reaction], bool]], + timeout: Optional[float] = None, + ) -> Reaction: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_reaction_add', 'raw_reaction_remove'], + /, + *, + check: Optional[Callable[[RawReactionActionEvent], bool]], + timeout: Optional[float] = None, + ) -> RawReactionActionEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_reaction_clear'], + /, + *, + check: Optional[Callable[[RawReactionClearEvent], bool]], + timeout: Optional[float] = None, + ) -> RawReactionClearEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_reaction_clear_emoji'], + /, + *, + check: Optional[Callable[[RawReactionClearEmojiEvent], bool]], + timeout: Optional[float] = None, + ) -> RawReactionClearEmojiEvent: + ... + + # Roles + + @overload + async def wait_for( + self, + event: Literal['guild_role_create', 'guild_role_delete'], + /, + *, + check: Optional[Callable[[Role], bool]], + timeout: Optional[float] = None, + ) -> Role: + ... + + @overload + async def wait_for( + self, + event: Literal['guild_role_update'], + /, + *, + check: Optional[Callable[[Role, Role], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Role, Role]: + ... + + # Scheduled Events + + @overload + async def wait_for( + self, + event: Literal['scheduled_event_create', 'scheduled_event_delete'], + /, + *, + check: Optional[Callable[[ScheduledEvent], bool]], + timeout: Optional[float] = None, + ) -> ScheduledEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['scheduled_event_user_add', 'scheduled_event_user_remove'], + /, + *, + check: Optional[Callable[[ScheduledEvent, User], bool]], + timeout: Optional[float] = None, + ) -> Tuple[ScheduledEvent, User]: + ... + + # Stages + + @overload + async def wait_for( + self, + event: Literal['stage_instance_create', 'stage_instance_delete'], + /, + *, + check: Optional[Callable[[StageInstance], bool]], + timeout: Optional[float] = None, + ) -> StageInstance: + ... + + @overload + async def wait_for( + self, + event: Literal['stage_instance_update'], + /, + *, + check: Optional[Callable[[StageInstance, StageInstance], bool]], + timeout: Optional[float] = None, + ) -> Coroutine[Any, Any, Tuple[StageInstance, StageInstance]]: + ... + + # Threads + @overload + async def wait_for( + self, + event: Literal['thread_create', 'thread_join', 'thread_remove', 'thread_delete'], + /, + *, + check: Optional[Callable[[Thread], bool]], + timeout: Optional[float] = None, + ) -> Thread: + ... + + @overload + async def wait_for( + self, + event: Literal['thread_update'], + /, + *, + check: Optional[Callable[[Thread, Thread], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Thread, Thread]: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_thread_update'], + /, + *, + check: Optional[Callable[[RawThreadUpdateEvent], bool]], + timeout: Optional[float] = None, + ) -> RawThreadUpdateEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_thread_delete'], + /, + *, + check: Optional[Callable[[RawThreadDeleteEvent], bool]], + timeout: Optional[float] = None, + ) -> RawThreadDeleteEvent: + ... + + @overload + async def wait_for( + self, + event: Literal['thread_member_join', 'thread_member_remove'], + /, + *, + check: Optional[Callable[[ThreadMember], bool]], + timeout: Optional[float] = None, + ) -> ThreadMember: + ... + + @overload + async def wait_for( + self, + event: Literal['raw_thread_member_remove'], + /, + *, + check: Optional[Callable[[RawThreadMembersUpdate], bool]], + timeout: Optional[float] = None, + ) -> RawThreadMembersUpdate: + ... + + # Voice + + @overload + async def wait_for( + self, + event: Literal['voice_state_update'], + /, + *, + check: Optional[Callable[[Member, VoiceState, VoiceState], bool]], + timeout: Optional[float] = None, + ) -> Tuple[Member, VoiceState, VoiceState]: + ... + + # Commands + + @overload + async def wait_for( + self: Union[Bot, AutoShardedBot], + event: Literal["command", "command_completion"], + /, + *, + check: Optional[Callable[[Context], bool]] = None, + timeout: Optional[float] = None, + ) -> Context: + ... + + @overload + async def wait_for( + self: Union[Bot, AutoShardedBot], + event: Literal["command_error"], + /, + *, + check: Optional[Callable[[Context, CommandError], bool]] = None, + timeout: Optional[float] = None, + ) -> Tuple[Context, CommandError]: + ... + + @overload + async def wait_for( self, event: str, /, @@ -1118,6 +1829,16 @@ class Client: check: Optional[Callable[..., bool]] = None, timeout: Optional[float] = None, ) -> Any: + ... + + def wait_for( + self, + event: str, + /, + *, + check: Optional[Callable[..., bool]] = None, + timeout: Optional[float] = None, + ) -> Coro[Any]: """|coro| Waits for a WebSocket event to be dispatched. @@ -1223,7 +1944,7 @@ class Client: # event registration - def event(self, coro: Coro, /) -> Coro: + def event(self, coro: CoroT, /) -> CoroT: """A decorator that registers an event to listen to. You can find more info about the events on the :ref:`documentation below `.