Fix typing issues and improve typing completeness across the library

Co-authored-by: Danny <Rapptz@users.noreply.github.com>
Co-authored-by: Josh <josh.ja.butt@gmail.com>

discord/ext/commands/bot.py

@@ -33,7 +33,21 @@ import importlib.util
 import sys
 import traceback
 import types
-from typing import Any, Callable, Mapping, List, Dict, TYPE_CHECKING, Optional, TypeVar, Type, Union, overload
+from typing import (
+    Any,
+    Callable,
+    Mapping,
+    List,
+    Dict,
+    TYPE_CHECKING,
+    Optional,
+    TypeVar,
+    Type,
+    Union,
+    Iterable,
+    Collection,
+    overload,
+)
 
 import discord
 from discord import app_commands
@@ -55,10 +69,18 @@ if TYPE_CHECKING:
     from discord.message import Message
     from discord.abc import User, Snowflake
     from ._types import (
+        _Bot,
+        BotT,
         Check,
         CoroFunc,
+        ContextT,
+        MaybeCoroFunc,
     )
 
+    _Prefix = Union[Iterable[str], str]
+    _PrefixCallable = MaybeCoroFunc[[BotT, Message], _Prefix]
+    PrefixType = Union[_Prefix, _PrefixCallable[BotT]]
+
 __all__ = (
     'when_mentioned',
     'when_mentioned_or',
@@ -68,11 +90,9 @@ __all__ = (
 
 T = TypeVar('T')
 CFT = TypeVar('CFT', bound='CoroFunc')
-CXT = TypeVar('CXT', bound='Context')
-BT = TypeVar('BT', bound='Union[Bot, AutoShardedBot]')
 
 
-def when_mentioned(bot: Union[Bot, AutoShardedBot], msg: Message) -> List[str]:
+def when_mentioned(bot: _Bot, msg: Message) -> List[str]:
     """A callable that implements a command prefix equivalent to being mentioned.
 
     These are meant to be passed into the :attr:`.Bot.command_prefix` attribute.
@@ -81,7 +101,7 @@ def when_mentioned(bot: Union[Bot, AutoShardedBot], msg: Message) -> List[str]:
     return [f'<@{bot.user.id}> ', f'<@!{bot.user.id}> ']  # type: ignore
 
 
-def when_mentioned_or(*prefixes: str) -> Callable[[Union[Bot, AutoShardedBot], Message], List[str]]:
+def when_mentioned_or(*prefixes: str) -> Callable[[_Bot, Message], List[str]]:
     """A callable that implements when mentioned or other prefixes provided.
 
     These are meant to be passed into the :attr:`.Bot.command_prefix` attribute.
@@ -124,27 +144,33 @@ class _DefaultRepr:
         return '<default-help-command>'
 
 
-_default = _DefaultRepr()
+_default: Any = _DefaultRepr()
 
 
-class BotBase(GroupMixin):
-    def __init__(self, command_prefix, help_command=_default, description=None, **options):
+class BotBase(GroupMixin[None]):
+    def __init__(
+        self,
+        command_prefix: PrefixType[BotT],
+        help_command: HelpCommand = _default,
+        description: Optional[str] = None,
+        **options: Any,
+    ) -> None:
         super().__init__(**options)
-        self.command_prefix = command_prefix
+        self.command_prefix: PrefixType[BotT] = command_prefix
         self.extra_events: Dict[str, List[CoroFunc]] = {}
         # Self doesn't have the ClientT bound, but since this is a mixin it technically does
         self.__tree: app_commands.CommandTree[Self] = app_commands.CommandTree(self)  # type: ignore
         self.__cogs: Dict[str, Cog] = {}
         self.__extensions: Dict[str, types.ModuleType] = {}
         self._checks: List[Check] = []
-        self._check_once = []
-        self._before_invoke = None
-        self._after_invoke = None
-        self._help_command = None
-        self.description = inspect.cleandoc(description) if description else ''
-        self.owner_id = options.get('owner_id')
-        self.owner_ids = options.get('owner_ids', set())
-        self.strip_after_prefix = options.get('strip_after_prefix', False)
+        self._check_once: List[Check] = []
+        self._before_invoke: Optional[CoroFunc] = None
+        self._after_invoke: Optional[CoroFunc] = None
+        self._help_command: Optional[HelpCommand] = None
+        self.description: str = inspect.cleandoc(description) if description else ''
+        self.owner_id: Optional[int] = options.get('owner_id')
+        self.owner_ids: Optional[Collection[int]] = options.get('owner_ids', set())
+        self.strip_after_prefix: bool = options.get('strip_after_prefix', False)
 
         if self.owner_id and self.owner_ids:
             raise TypeError('Both owner_id and owner_ids are set.')
@@ -182,7 +208,7 @@ class BotBase(GroupMixin):
 
         await super().close()  # type: ignore
 
-    async def on_command_error(self, context: Context, exception: errors.CommandError) -> None:
+    async def on_command_error(self, context: Context[BotT], exception: errors.CommandError) -> None:
         """|coro|
 
         The default command error handler provided by the bot.
@@ -237,7 +263,7 @@ class BotBase(GroupMixin):
         self.add_check(func)  # type: ignore
         return func
 
-    def add_check(self, func: Check, /, *, call_once: bool = False) -> None:
+    def add_check(self, func: Check[ContextT], /, *, call_once: bool = False) -> None:
         """Adds a global check to the bot.
 
         This is the non-decorator interface to :meth:`.check`
@@ -261,7 +287,7 @@ class BotBase(GroupMixin):
         else:
             self._checks.append(func)
 
-    def remove_check(self, func: Check, /, *, call_once: bool = False) -> None:
+    def remove_check(self, func: Check[ContextT], /, *, call_once: bool = False) -> None:
         """Removes a global check from the bot.
 
         This function is idempotent and will not raise an exception
@@ -324,7 +350,7 @@ class BotBase(GroupMixin):
         self.add_check(func, call_once=True)
         return func
 
-    async def can_run(self, ctx: Context, *, call_once: bool = False) -> bool:
+    async def can_run(self, ctx: Context[BotT], *, call_once: bool = False) -> bool:
         data = self._check_once if call_once else self._checks
 
         if len(data) == 0:
@@ -947,7 +973,7 @@ class BotBase(GroupMixin):
             # if the load failed, the remnants should have been
             # cleaned from the load_extension function call
             # so let's load it from our old compiled library.
-            await lib.setup(self)  # type: ignore
+            await lib.setup(self)
             self.__extensions[name] = lib
 
             # revert sys.modules back to normal and raise back to caller
@@ -1015,11 +1041,12 @@ class BotBase(GroupMixin):
         """
         prefix = ret = self.command_prefix
         if callable(prefix):
-            ret = await discord.utils.maybe_coroutine(prefix, self, message)
+            # self will be a Bot or AutoShardedBot
+            ret = await discord.utils.maybe_coroutine(prefix, self, message)  # type: ignore
 
         if not isinstance(ret, str):
             try:
-                ret = list(ret)
+                ret = list(ret)  # type: ignore
             except TypeError:
                 # It's possible that a generator raised this exception.  Don't
                 # replace it with our own error if that's the case.
@@ -1048,15 +1075,15 @@ class BotBase(GroupMixin):
         self,
         message: Message,
         *,
-        cls: Type[CXT] = ...,
-    ) -> CXT:  # type: ignore
+        cls: Type[ContextT] = ...,
+    ) -> ContextT:
         ...
 
     async def get_context(
         self,
         message: Message,
         *,
-        cls: Type[CXT] = MISSING,
+        cls: Type[ContextT] = MISSING,
     ) -> Any:
         r"""|coro|
 
@@ -1137,7 +1164,7 @@ class BotBase(GroupMixin):
         ctx.command = self.all_commands.get(invoker)
         return ctx
 
-    async def invoke(self, ctx: Context) -> None:
+    async def invoke(self, ctx: Context[BotT]) -> None:
         """|coro|
 
         Invokes the command given under the invocation context and
@@ -1189,9 +1216,10 @@ class BotBase(GroupMixin):
             return
 
         ctx = await self.get_context(message)
-        await self.invoke(ctx)
+        # the type of the invocation context's bot attribute will be correct
+        await self.invoke(ctx)  # type: ignore
 
-    async def on_message(self, message):
+    async def on_message(self, message: Message) -> None:
         await self.process_commands(message)