Type checking fixes

dozer/bot.py

@@ -1,14 +1,14 @@
 """Bot object for Dozer"""
-
-
 import os
 import re
 import sys
 import traceback
-from typing import Pattern
+from typing import Pattern, Optional, Union, Generator, Dict, Any
 
 import discord
+from discord import Status, Message
 from discord.ext import commands
+from discord.ext.commands import Cooldown, CommandError, BucketType
 from loguru import logger
 from sentry_sdk import capture_exception
 
@@ -18,6 +18,7 @@
 from .context import DozerContext
 from .db import db_init, db_migrate
 
+
 if discord.version_info.major < 2:
     logger.error("Your installed discord.py version is too low "
                  "%d.%d.%d, please upgrade to at least 2.0.0",
@@ -36,14 +37,14 @@ class InvalidContext(commands.CheckFailure):
 
 class Dozer(commands.Bot):
     """Botty things that are critical to Dozer working"""
-    _global_cooldown = commands.Cooldown(1, 1)  # One command per second per user
+    _global_cooldown: Cooldown = Cooldown(1, 1)  # One command per second per user
 
-    def __init__(self, config: dict, *args, **kwargs):
+    def __init__(self, config: Dict[str, Union[Dict[str, str], str]], *args, **kwargs):
         self.wavelink = None
-        self.dynamic_prefix = _utils.PrefixHandler(config['prefix'])
+        self.dynamic_prefix: _utils.PrefixHandler = _utils.PrefixHandler(str(config['prefix']))
         super().__init__(command_prefix=self.dynamic_prefix.handler, *args, **kwargs)
-        self.config = config
-        self._restarting = False
+        self.config: Dict[str, Any] = config
+        self._restarting: bool = False
         self.check(self.global_checks)
 
     async def setup_hook(self) -> None:
@@ -67,23 +68,24 @@ async def on_ready(self):
                 perms |= cmd.required_permissions.value
             else:
                 logger.warning(f"Command {cmd} not subclass of Dozer type.")
-        logger.debug('Bot Invite: {}'.format(utils.oauth_url(self.user.id, discord.Permissions(perms))))
+        logger.debug('Bot Invite: {}'.format(utils.oauth_url(str(self.user.id), discord.Permissions(perms))))
         if self.config['is_backup']:
-            status = discord.Status.dnd
+            status: Status = Status.dnd
         else:
-            status = discord.Status.online
-        activity = discord.Game(name=f"@{self.user.name} or '{self.config['prefix']}' in {len(self.guilds)} guilds")
+            status: Status = Status.online
+        activity: discord.Game = discord.Game(name=f"@{self.user.name} or '{self.config['prefix']}' in {len(self.guilds)} guilds")
         try:
             await self.change_presence(activity=activity, status=status)
         except TypeError:
             logger.warning("You are running an older version of the discord.py rewrite (with breaking changes)! "
                            "To upgrade, run `pip install -r requirements.txt --upgrade`")
 
-    async def get_context(self, message: discord.Message, *, cls=DozerContext):  # pylint: disable=arguments-differ
+    async def get_context(self, message: Message, *, cls=DozerContext) -> DozerContext:  # pylint: disable=arguments-differ
         ctx = await super().get_context(message, cls=cls)
+        ctx.prefix = self.dynamic_prefix.handler(self, message)
         return ctx
 
-    async def on_command_error(self, context: DozerContext, exception):  # pylint: disable=arguments-differ
+    async def on_command_error(self, context: DozerContext, exception: CommandError):  # pylint: disable=arguments-differ
         if isinstance(exception, commands.NoPrivateMessage):
             await context.send('{}, This command cannot be used in DMs.'.format(context.author.mention))
         elif isinstance(exception, commands.UserInputError):
@@ -105,11 +107,14 @@ async def on_command_error(self, context: DozerContext, exception):  # pylint: d
                 '{}, That command is on cooldown! Try again in {:.2f}s!'.format(context.author.mention,
                                                                                 exception.retry_after))
         elif isinstance(exception, commands.MaxConcurrencyReached):
-            types = {discord.ext.commands.BucketType.default: "`Global`",
-                     discord.ext.commands.BucketType.guild: "`Guild`",
-                     discord.ext.commands.BucketType.channel: "`Channel`",
-                     discord.ext.commands.BucketType.category: "`Category`",
-                     discord.ext.commands.BucketType.member: "`Member`", discord.ext.commands.BucketType.user: "`User`"}
+            types: Dict[BucketType, str] = {
+                BucketType.default: "`Global`",
+                BucketType.guild: "`Guild`",
+                BucketType.channel: "`Channel`",
+                BucketType.category: "`Category`",
+                BucketType.member: "`Member`",
+                BucketType.user: "`User`"
+            }
             await context.send(
                 '{}, That command has exceeded the max {} concurrency limit of `{}` instance! Please try again later.'.format(
                     context.author.mention, types[exception.per], exception.number))
@@ -128,14 +133,14 @@ async def on_command_error(self, context: DozerContext, exception):  # pylint: d
                              context.channel.recipient, context.message.content)
             logger.error(''.join(traceback.format_exception(type(exception), exception, exception.__traceback__)))
 
-    async def on_error(self, event_method, *args, **kwargs):
+    async def on_error(self, event_method: str, *args, **kwargs):
         """Don't ignore the error, causing Sentry to capture it."""
         print('Ignoring exception in {}'.format(event_method), file=sys.stderr)
         traceback.print_exc()
         capture_exception()
 
     @staticmethod
-    def format_error(ctx: DozerContext, err: Exception, *, word_re: Pattern = re.compile('[A-Z][a-z]+')):
+    def format_error(ctx: DozerContext, err: Exception, *, word_re: Pattern = re.compile('[A-Z][a-z]+')) -> str:
         """Turns an exception into a user-friendly (or -friendlier, at least) error message."""
         type_words = word_re.findall(type(err).__name__)
         type_msg = ' '.join(map(str.lower, type_words))
@@ -145,7 +150,7 @@ def format_error(ctx: DozerContext, err: Exception, *, word_re: Pattern = re.com
         else:
             return type_msg
 
-    def global_checks(self, ctx: DozerContext):
+    def global_checks(self, ctx: DozerContext) -> bool:
         """Checks that should be executed before passed to the command"""
         if ctx.author.bot:
             raise InvalidContext('Bots cannot run commands!')
@@ -154,6 +159,15 @@ def global_checks(self, ctx: DozerContext):
             raise InvalidContext('Global rate-limit exceeded!')
         return True
 
+    def get_command(self, name: str) -> Optional[Union[_utils.Command, _utils.Group]]:  # pylint: disable=arguments-differ
+        return super().get_command(name)
+
+    def walk_commands(self) -> Generator[Union[_utils.Command, _utils.Group], None, None]:
+        return super().walk_commands()
+
+    def get_cog(self, name: str, /) -> Optional[_utils.Cog]:
+        return super().get_cog(name)
+
     def run(self, *args, **kwargs):
         token = self.config['discord_token']
         del self.config['discord_token']  # Prevent token dumping