Use an enum for callbacks

Author: CraftSpider

discord/ext/test/_types.py

@@ -8,32 +8,11 @@
 import discord
 import typing
 
-if typing.TYPE_CHECKING:
-    from discord.types import (
-        role, gateway, appinfo, user, guild, emoji, channel, message, sticker,  # noqa: F401
-        scheduled_event, member  # noqa: F401
-    )
-
-    AnyChannelJson = channel.VoiceChannel | channel.TextChannel | channel.DMChannel | channel.CategoryChannel
-else:
-    class OpenNamespace:
-        def __getattr__(self, item: str) -> Self:
-            return self
-
-        def __subclasscheck__(self, subclass: type) -> Literal[True]:
-            return True
-
-        def __or__(self, other: T) -> T:
-            return other
-
-    def __getattr__(name: str) -> OpenNamespace:
-        return OpenNamespace()
-
 T = TypeVar('T')
 P = ParamSpec('P')
 
-Callback = Callable[P, Coroutine[None, None, None]]
-AnyChannel = (discord.abc.GuildChannel | discord.TextChannel | discord.VoiceChannel | discord.StageChannel | discord.DMChannel | discord.Thread | discord.GroupChannel)
+AnyChannel = (discord.abc.GuildChannel | discord.TextChannel | discord.VoiceChannel | discord.StageChannel
+              | discord.DMChannel | discord.Thread | discord.GroupChannel)
 
 
 class Wrapper(Protocol[P, T]):
@@ -55,3 +34,25 @@ class Undef(Enum):
 
 
 undefined: Literal[Undef.undefined] = Undef.undefined
+
+
+if typing.TYPE_CHECKING:
+    from discord.types import (
+        role, gateway, appinfo, user, guild, emoji, channel, message, sticker, snowflake,  # noqa: F401
+        scheduled_event, member  # noqa: F401
+    )
+
+    AnyChannelJson = channel.VoiceChannel | channel.TextChannel | channel.DMChannel | channel.CategoryChannel
+else:
+    class OpenNamespace:
+        def __getattr__(self, item: str) -> Self:
+            return self
+
+        def __subclasscheck__(self, subclass: type) -> Literal[True]:
+            return True
+
+        def __or__(self, other: T) -> T:
+            return other
+
+    def __getattr__(name: str) -> OpenNamespace:
+        return OpenNamespace()

discord/ext/test/backend.py

@@ -26,6 +26,8 @@
 from ._types import Undef, undefined
 from discord.types.snowflake import Snowflake
 
+from .callbacks import CallbackEvent
+
 
 class BackendState(NamedTuple):
     """
@@ -140,7 +142,7 @@ async def delete_channel(self, channel_id: Snowflake, *, reason: str | None = No
             delete_channel(channel)
 
     async def get_channel(self, channel_id: Snowflake) -> _types.channel.GuildChannel:
-        await callbacks.dispatch_event("get_channel", channel_id)
+        await callbacks.dispatch_event(CallbackEvent.get_channel, channel_id)
 
         find = None
         for guild in get_state().guilds:
@@ -155,7 +157,7 @@ async def start_private_message(self, user_id: Snowflake) -> _types.channel.DMCh
         locs = _get_higher_locs(1)
         user = locs["self"]
 
-        await callbacks.dispatch_event("start_private_message", user)
+        await callbacks.dispatch_event(CallbackEvent.start_private_message, user)
 
         return facts.make_dm_channel_dict(user)
 
@@ -212,22 +214,22 @@ async def send_message(
                                attachments=attachments,
                                nonce=nonce
                                )
-        await callbacks.dispatch_event("send_message", message)
+        await callbacks.dispatch_event(CallbackEvent.send_message, message)
 
         return facts.dict_from_object(message)
 
     async def send_typing(self, channel_id: Snowflake) -> None:
         locs = _get_higher_locs(1)
         channel = locs.get("channel", None)
 
-        await callbacks.dispatch_event("send_typing", channel)
+        await callbacks.dispatch_event(CallbackEvent.send_typing, channel)
 
     async def delete_message(self, channel_id: Snowflake, message_id: Snowflake, *,
                              reason: str | None = None) -> None:
         locs = _get_higher_locs(1)
         message = locs["self"]
 
-        await callbacks.dispatch_event("delete_message", message.channel, message, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.delete_message, message.channel, message, reason=reason)
 
         delete_message(message)
 
@@ -236,7 +238,7 @@ async def edit_message(self, channel_id: Snowflake, message_id: Snowflake,
         locs = _get_higher_locs(1)
         message = locs["self"]
 
-        await callbacks.dispatch_event("edit_message", message.channel, message, fields)
+        await callbacks.dispatch_event(CallbackEvent.edit_message, message.channel, message, fields)
 
         return edit_message(message, **fields)
 
@@ -250,7 +252,7 @@ async def add_reaction(self, channel_id: Snowflake, message_id: Snowflake,
 
         emoji = emoji  # TODO: Turn this back into class?
 
-        await callbacks.dispatch_event("add_reaction", message, emoji)
+        await callbacks.dispatch_event(CallbackEvent.add_reaction, message, emoji)
 
         add_reaction(message, user, emoji)
 
@@ -261,7 +263,7 @@ async def remove_reaction(self, channel_id: Snowflake, message_id: Snowflake,
         message = locs["self"]
         member = locs["member"]
 
-        await callbacks.dispatch_event("remove_reaction", message, emoji, member)
+        await callbacks.dispatch_event(CallbackEvent.remove_reaction, message, emoji, member)
 
         remove_reaction(message, member, emoji)
 
@@ -271,7 +273,7 @@ async def remove_own_reaction(self, channel_id: Snowflake, message_id: Snowflake
         message = locs["self"]
         member = locs["member"]
 
-        await callbacks.dispatch_event("remove_own_reaction", message, emoji, member)
+        await callbacks.dispatch_event(CallbackEvent.remove_own_reaction, message, emoji, member)
 
         remove_reaction(message, self.state.user, emoji)
 
@@ -285,7 +287,7 @@ async def get_message(self, channel_id: Snowflake,
         locs = _get_higher_locs(1)
         channel = locs["self"]
 
-        await callbacks.dispatch_event("get_message", channel, message_id)
+        await callbacks.dispatch_event(CallbackEvent.get_message, channel, message_id)
 
         messages = get_config().messages[int(channel_id)]
         find = next(filter(lambda m: m["id"] == message_id, messages), None)
@@ -304,7 +306,7 @@ async def logs_from(
         locs = _get_higher_locs(1)
         channel = locs["self"]
 
-        await callbacks.dispatch_event("logs_from", channel, limit, before=None, after=None, around=None)
+        await callbacks.dispatch_event(CallbackEvent.logs_from, channel, limit, before=None, after=None, around=None)
 
         messages = get_config().messages[int(channel_id)]
         if after is not None:
@@ -326,7 +328,7 @@ async def kick(self, user_id: Snowflake, guild_id: Snowflake,
         guild = locs["self"]
         member = locs["user"]
 
-        await callbacks.dispatch_event("kick", guild, member, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.kick, guild, member, reason=reason)
 
         delete_member(member)
 
@@ -337,7 +339,7 @@ async def ban(self, user_id: Snowflake, guild_id: Snowflake,
         guild = locs["self"]
         member = locs["user"]
 
-        await callbacks.dispatch_event("ban", guild, member, delete_message_days, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.ban, guild, member, delete_message_days, reason=reason)
 
         delete_member(member)
 
@@ -346,7 +348,7 @@ async def unban(self, user_id: Snowflake, guild_id: Snowflake, *,
         locs = _get_higher_locs(1)
         guild = locs["self"]
         member = locs["user"]
-        await callbacks.dispatch_event("unban", guild, member, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.unban, guild, member, reason=reason)
 
     async def change_my_nickname(self, guild_id: Snowflake, nickname: str, *,
                                  reason: str | None = None) -> _types.member.Nickname:
@@ -355,7 +357,7 @@ async def change_my_nickname(self, guild_id: Snowflake, nickname: str, *,
 
         me.nick = nickname
 
-        await callbacks.dispatch_event("change_nickname", nickname, me, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.change_nickname, nickname, me, reason=reason)
 
         return {"nick": nickname}
 
@@ -365,7 +367,7 @@ async def edit_member(self, guild_id: Snowflake, user_id: Snowflake, *,
         locs = _get_higher_locs(1)
         member = locs["self"]
 
-        await callbacks.dispatch_event("edit_member", fields, member, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.edit_member, fields, member, reason=reason)
         member = update_member(member, nick=fields.get('nick'), roles=fields.get('roles'))
         return facts.dict_from_object(member)
 
@@ -384,7 +386,7 @@ async def edit_role(self, guild_id: Snowflake, role_id: Snowflake, *,
         role = locs["self"]
         guild = role.guild
 
-        await callbacks.dispatch_event("edit_role", guild, role, fields, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.edit_role, guild, role, fields, reason=reason)
 
         update_role(role, **fields)
         return facts.dict_from_object(role)
@@ -395,7 +397,7 @@ async def delete_role(self, guild_id: Snowflake, role_id: Snowflake, *,
         role = locs["self"]
         guild = role.guild
 
-        await callbacks.dispatch_event("delete_role", guild, role, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.delete_role, guild, role, reason=reason)
 
         delete_role(role)
 
@@ -405,7 +407,7 @@ async def create_role(self, guild_id: Snowflake, *, reason: str | None = None,
         guild = locs["self"]
         role = make_role(guild=guild, **fields)
 
-        await callbacks.dispatch_event("create_role", guild, role, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.create_role, guild, role, reason=reason)
 
         return facts.dict_from_object(role)
 
@@ -416,7 +418,7 @@ async def move_role_position(self, guild_id: Snowflake,
         role = locs["self"]
         guild = role.guild
 
-        await callbacks.dispatch_event("move_role", guild, role, positions, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.move_role, guild, role, positions, reason=reason)
 
         for pair in positions:
             guild._roles[pair["id"]].position = pair["position"]
@@ -428,7 +430,7 @@ async def add_role(self, guild_id: Snowflake, user_id: Snowflake,
         member = locs["self"]
         role = locs["role"]
 
-        await callbacks.dispatch_event("add_role", member, role, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.add_role, member, role, reason=reason)
 
         roles = [role] + [x for x in member.roles if x.id != member.guild.id]
         update_member(member, roles=roles)
@@ -440,7 +442,7 @@ async def remove_role(self, guild_id: Snowflake, user_id: Snowflake,
         member = locs["self"]
         role = locs["role"]
 
-        await callbacks.dispatch_event("remove_role", member, role, reason=reason)
+        await callbacks.dispatch_event(CallbackEvent.remove_role, member, role, reason=reason)
 
         roles = [x for x in member.roles if x != role and x.id != member.guild.id]
         update_member(member, roles=roles)
@@ -463,7 +465,7 @@ async def application_info(self) -> _types.appinfo.AppInfo:
         }
 
         appinfo = discord.AppInfo(self.state, data)
-        await callbacks.dispatch_event("app_info", appinfo)
+        await callbacks.dispatch_event(CallbackEvent.app_info, appinfo)
 
         return data
 
@@ -540,7 +542,13 @@ async def get_guilds(self, limit: int, before: Snowflake | None = None,
                          after: Snowflake | None = None,
                          with_counts: bool = True) -> list[_types.guild.Guild]:
         # self.request(Route('GET', '/users/@me/guilds')
-        await callbacks.dispatch_event("get_guilds", limit, before=before, after=after, with_counts=with_counts)
+        await callbacks.dispatch_event(
+            CallbackEvent.get_guilds,
+            limit,
+            before=before,
+            after=after,
+            with_counts=with_counts,
+        )
         guilds = get_state().guilds  # List[]
 
         guilds_new = [facts.dict_from_object(guild) for guild in guilds]

discord/ext/test/callbacks.py

@@ -6,14 +6,52 @@
 
 import logging
 import typing
+import discord
+from enum import Enum
+from typing import Callable, overload, Literal, Any, Awaitable
+
 from . import _types
 
+GetChannelCallback = Callable[[_types.snowflake.Snowflake], Awaitable[None]]
+SendMessageCallback = Callable[[discord.Message], Awaitable[None]]
+EditMemberCallback = Callable[[dict[str, Any], discord.Member, str | None], Awaitable[None]]
+Callback = GetChannelCallback | SendMessageCallback | EditMemberCallback | Callable[..., Awaitable[None]]
+
 log = logging.getLogger("discord.ext.tests")
 
-_callbacks: dict[str, _types.Callback[...]] = {}
+
+class CallbackEvent(Enum):
+    get_channel = "get_channel"
+    presence = "presence"
+    start_private_message = "start_private_message"
+    send_message = "send_message"
+    send_typing = "send_typing"
+    delete_message = "delete_message"
+    edit_message = "edit_message"
+    add_reaction = "add_reaction"
+    remove_reaction = "remove_reaction"
+    remove_own_reaction = "remove_own_reaction"
+    get_message = "get_message"
+    logs_from = "logs_from"
+    kick = "kick"
+    ban = "ban"
+    unban = "unban"
+    change_nickname = "change_nickname"
+    edit_member = "edit_member"
+    create_role = "create_role"
+    edit_role = "edit_role"
+    delete_role = "delete_role"
+    move_role = "move_role"
+    add_role = "add_role"
+    remove_role = "remove_role"
+    app_info = "app_info"
+    get_guilds = "get_guilds"
 
 
-async def dispatch_event(event: str, *args: typing.Any, **kwargs: typing.Any) -> None:
+_callbacks: dict[CallbackEvent, Callback] = {}
+
+
+async def dispatch_event(event: CallbackEvent, *args: typing.Any, **kwargs: typing.Any) -> None:
     """
         Dispatch an event to a set handler, if one exists. Will ignore handler errors,
         just print a log
@@ -30,7 +68,19 @@ async def dispatch_event(event: str, *args: typing.Any, **kwargs: typing.Any) ->
             log.error(f"Error in handler for event {event}: {e}")
 
 
-def set_callback(cb: _types.Callback[...], event: str) -> None:
+@overload
+def set_callback(cb: GetChannelCallback, event: Literal[CallbackEvent.get_channel]) -> None: ...
+
+
+@overload
+def set_callback(cb: SendMessageCallback, event: Literal[CallbackEvent.send_message]) -> None: ...
+
+
+@overload
+def set_callback(cb: EditMemberCallback, event: Literal[CallbackEvent.edit_member]) -> None: ...
+
+
+def set_callback(cb: Callback, event: CallbackEvent) -> None:
     """
         Set the callback to use for a specific event
 
@@ -40,7 +90,7 @@ def set_callback(cb: _types.Callback[...], event: str) -> None:
     _callbacks[event] = cb
 
 
-def get_callback(event: str) -> _types.Callback[...]:
+def get_callback(event: CallbackEvent) -> Callback:
     """
         Get the current callback for an event, or raise an exception if one isn't set
 
@@ -52,7 +102,7 @@ def get_callback(event: str) -> _types.Callback[...]:
     return _callbacks[event]
 
 
-def remove_callback(event: str) -> _types.Callback[...] | None:
+def remove_callback(event: CallbackEvent) -> Callback | None:
     """
         Remove the callback set for an event, returning it, or None if one isn't set