Add overloads for all the events in the event reference

Co-authored-by: Sebastian Law <sebnlaw@gmail.com>
Co-authored-by: Josh <8677174+bijij@users.noreply.github.com>

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 <discord-api-events>`.