Type checking

dozer/__main__.py

@@ -1,17 +1,15 @@
 """Initializes the bot and deals with the configuration file"""
 
+import asyncio
 import json
 import os
 import sys
-import asyncio
 
 import discord
 import sentry_sdk
 
 from .db import db_init, db_migrate
 
-from . import db
-
 config = {
     'prefix': '&', 'developers': [],
     'cache_size': 20000,
@@ -61,7 +59,7 @@
     json.dump(config, f, indent='\t')
 
 if config['sentry_url'] != "":
-    sentry_sdk.init(
+    sentry_sdk.init(  # pylint: disable=abstract-class-instantiated  # noqa: E0110
         config['sentry_url'],
         traces_sample_rate=1.0,
     )
@@ -82,7 +80,6 @@
 
 bot = Dozer(config, intents=intents, max_messages=config['cache_size'])
 
-
 for ext in os.listdir('dozer/cogs'):
     if not ext.startswith(('_', '.')):
         bot.load_extension('dozer.cogs.' + ext[:-3])  # Remove '.py'

dozer/bot.py

@@ -4,6 +4,7 @@
 import re
 import sys
 import traceback
+from typing import Pattern
 
 import discord
 from discord.ext import commands
@@ -12,6 +13,7 @@
 
 from . import utils
 from .cogs import _utils
+from .context import DozerContext
 
 DOZER_LOGGER = logging.getLogger('dozer')
 DOZER_LOGGER.level = logging.INFO
@@ -36,19 +38,11 @@ class InvalidContext(commands.CheckFailure):
     """
 
 
-class DozerContext(commands.Context):
-    """Cleans all messages before sending"""
-    async def send(self, content=None, **kwargs): # pylint: disable=arguments-differ
-        if content is not None:
-            content = utils.clean(self, content, mass=True, member=False, role=False, channel=False)
-        return await super().send(content, **kwargs)
-
-
 class Dozer(commands.Bot):
     """Botty things that are critical to Dozer working"""
     _global_cooldown = commands.Cooldown(1, 1, commands.BucketType.user)  # One command per second per user
 
-    def __init__(self, config, *args, **kwargs):
+    def __init__(self, config: dict, *args, **kwargs):
         self.dynamic_prefix = _utils.PrefixHandler(config['prefix'])
         super().__init__(command_prefix=self.dynamic_prefix.handler, *args, **kwargs)
         self.slash = SlashCommand(self, sync_commands=True, override_type=True)
@@ -78,45 +72,52 @@ async def on_ready(self):
             DOZER_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, *, cls=DozerContext):
+    async def get_context(self, message: discord.Message, *, cls=DozerContext):
         ctx = await super().get_context(message, cls=cls)
         return ctx
 
-    async def on_command_error(self, context, exception):
+    async def on_command_error(self, context: DozerContext, exception):
         if isinstance(exception, commands.NoPrivateMessage):
             await context.send('{}, This command cannot be used in DMs.'.format(context.author.mention))
         elif isinstance(exception, commands.UserInputError):
             await context.send('{}, {}'.format(context.author.mention, self.format_error(context, exception)))
         elif isinstance(exception, commands.NotOwner):
             await context.send('{}, {}'.format(context.author.mention, exception.args[0]))
         elif isinstance(exception, commands.MissingPermissions):
-            permission_names = [name.replace('guild', 'server').replace('_', ' ').title() for name in exception.missing_perms]
+            permission_names = [name.replace('guild', 'server').replace('_', ' ').title() for name in
+                                exception.missing_perms]
             await context.send('{}, you need {} permissions to run this command!'.format(
                 context.author.mention, utils.pretty_concat(permission_names)))
         elif isinstance(exception, commands.BotMissingPermissions):
-            permission_names = [name.replace('guild', 'server').replace('_', ' ').title() for name in exception.missing_perms]
+            permission_names = [name.replace('guild', 'server').replace('_', ' ').title() for name in
+                                exception.missing_perms]
             await context.send('{}, I need {} permissions to run this command!'.format(
                 context.author.mention, utils.pretty_concat(permission_names)))
         elif isinstance(exception, commands.CommandOnCooldown):
             await context.send(
-                '{}, That command is on cooldown! Try again in {:.2f}s!'.format(context.author.mention, exception.retry_after))
+                '{}, 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`",
+            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`"}
             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))
         elif isinstance(exception, (commands.CommandNotFound, InvalidContext)):
             pass  # Silent ignore
         else:
-            await context.send('```\n%s\n```' % ''.join(traceback.format_exception_only(type(exception), exception)).strip())
+            await context.send(
+                '```\n%s\n```' % ''.join(traceback.format_exception_only(type(exception), exception)).strip())
             if isinstance(context.channel, discord.TextChannel):
                 DOZER_LOGGER.error('Error in command <%d> (%d.name!r(%d.id) %d(%d.id) %d(%d.id) %d)',
                                    context.command, context.guild, context.guild, context.channel, context.channel,
                                    context.author, context.author, context.message.content)
             else:
-                DOZER_LOGGER.error('Error in command <%d> (DM %d(%d.id) %d)', context.command, context.channel.recipient,
+                DOZER_LOGGER.error('Error in command <%d> (DM %d(%d.id) %d)', context.command,
+                                   context.channel.recipient,
                                    context.channel.recipient, context.message.content)
             DOZER_LOGGER.error(''.join(traceback.format_exception(type(exception), exception, exception.__traceback__)))
 
@@ -126,12 +127,12 @@ async def on_error(self, event_method, *args, **kwargs):
         traceback.print_exc()
         capture_exception()
 
-    async def on_slash_command_error(self, ctx, ex):
+    async def on_slash_command_error(self, ctx: DozerContext, ex: Exception):
         """Passes slash command errors to primary command handler"""
         await self.on_command_error(ctx, ex)
 
     @staticmethod
-    def format_error(ctx, err, *, word_re=re.compile('[A-Z][a-z]+')):
+    def format_error(ctx: DozerContext, err: Exception, *, word_re: Pattern = re.compile('[A-Z][a-z]+')):
         """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))
@@ -141,12 +142,12 @@ def format_error(ctx, err, *, word_re=re.compile('[A-Z][a-z]+')):
         else:
             return type_msg
 
-    def global_checks(self, ctx):
+    def global_checks(self, ctx: DozerContext):
         """Checks that should be executed before passed to the command"""
         if ctx.author.bot:
             raise InvalidContext('Bots cannot run commands!')
         retry_after = self._global_cooldown.update_rate_limit()
-        if retry_after and not hasattr(ctx, "is_pseudo"): # bypass ratelimit for su'ed commands
+        if retry_after and not hasattr(ctx, "is_pseudo"):  # bypass ratelimit for su'ed commands
             raise InvalidContext('Global rate-limit exceeded!')
         return True
 
@@ -155,7 +156,7 @@ def run(self, *args, **kwargs):
         del self.config['discord_token']  # Prevent token dumping
         super().run(token)
 
-    async def shutdown(self, restart=False):
+    async def shutdown(self, restart: bool = False):
         """Shuts down the bot"""
         self._restarting = restart
         await self.logout()

dozer/cogs/_utils.py

@@ -4,13 +4,16 @@
 import logging
 import typing
 from collections.abc import Mapping
+from typing import Dict
 
 import discord
 from discord.ext import commands
 
 from dozer import db
+from dozer.context import DozerContext
 
-__all__ = ['bot_has_permissions', 'command', 'group', 'Cog', 'Reactor', 'Paginator', 'paginate', 'chunk', 'dev_check', 'DynamicPrefixEntry']
+__all__ = ['bot_has_permissions', 'command', 'group', 'Cog', 'Reactor', 'Paginator', 'paginate', 'chunk', 'dev_check',
+           'DynamicPrefixEntry']
 
 DOZER_LOGGER = logging.getLogger("dozer")
 
@@ -81,15 +84,15 @@ def group(**kwargs):
 class Cog(commands.Cog):
     """Initiates cogs."""
 
-    def __init__(self, bot):
+    def __init__(self, bot: commands.Bot):
         super().__init__()
         self.bot = bot
 
 
 def dev_check():
     """Function decorator to check that the calling user is a developer"""
 
-    async def predicate(ctx):
+    async def predicate(ctx: DozerContext):
         if ctx.author.id not in ctx.bot.config['developers']:
             raise commands.NotOwner('you are not a developer!')
         return True
@@ -118,7 +121,7 @@ class Reactor:
     """
     _stop_reaction = object()
 
-    def __init__(self, ctx, initial_reactions, *, auto_remove=True, timeout=60):
+    def __init__(self, ctx: DozerContext, initial_reactions, *, auto_remove: bool = True, timeout: int = 60):
         """
         ctx: command context
         initial_reactions: iterable of emoji to react with on start
@@ -142,7 +145,8 @@ async def __aiter__(self):
             await self.message.add_reaction(emoji)
         while True:
             try:
-                reaction, reacting_member = await self.bot.wait_for('reaction_add', check=self._check_reaction, timeout=self.timeout)
+                reaction, reacting_member = await self.bot.wait_for('reaction_add', check=self._check_reaction,
+                                                                    timeout=self.timeout)
             except asyncio.TimeoutError:
                 break
 
@@ -167,7 +171,7 @@ def stop(self):
         """Listener for stop reactions."""
         self._action = self._stop_reaction
 
-    def _check_reaction(self, reaction, member):
+    def _check_reaction(self, reaction: discord.Reaction, member: discord.Member):
         return reaction.message.id == self.message.id and member.id == self.caller.id
 
 
@@ -195,7 +199,8 @@ class Paginator(Reactor):
         '\N{BLACK SQUARE FOR STOP}'  # :stop_button:
     )
 
-    def __init__(self, ctx, initial_reactions, pages, *, start=0, auto_remove=True, timeout=60):
+    def __init__(self, ctx: DozerContext, initial_reactions, pages, *, start: int = 0, auto_remove: bool = True,
+                 timeout: int = 60):
         all_reactions = list(initial_reactions)
         ind = all_reactions.index(Ellipsis)
         all_reactions[ind:ind + 1] = self.pagination_reactions
@@ -229,40 +234,41 @@ async def __aiter__(self):
                 else:  # Only valid option left is 4
                     self.stop()
 
-    def go_to_page(self, page):
+    def go_to_page(self, page: int):
         """Goes to a specific help page"""
         if isinstance(page, int):
             page = page % self.len_pages
             if page < 0:
                 page += self.len_pages
         self.page = page
-        self.do(self.message.edit(embed=self.pages[self.page]))
+        if self.message is not None:
+            self.do(self.message.edit(embed=self.pages[self.page]))
 
-    def next(self, amt=1):
+    def next(self, amt: int = 1):
         """Goes to the next help page"""
         if isinstance(self.page, int):
             self.go_to_page(self.page + amt)
         else:
             self.go_to_page(amt - 1)
 
-    def prev(self, amt=1):
+    def prev(self, amt: int = 1):
         """Goes to the previous help page"""
         if isinstance(self.page, int):
             self.go_to_page(self.page - amt)
         else:
             self.go_to_page(-amt)
 
 
-async def paginate(ctx, pages, *, start=0, auto_remove=True, timeout=60):
+async def paginate(ctx: DozerContext, pages, *, start: int = 0, auto_remove: bool = True, timeout: int = 60):
     """
     Simple pagination based on Paginator. Pagination is handled normally and other reactions are ignored.
     """
-    paginator = Paginator(ctx, (...,), pages, start=start, auto_remove=auto_remove, timeout=timeout)
+    paginator = Paginator(ctx, ..., pages, start=start, auto_remove=auto_remove, timeout=timeout)
     async for reaction in paginator:
         pass  # The normal pagination reactions are handled - just drop anything else
 
 
-def chunk(iterable, size):
+def chunk(iterable, size: int):
     """
     Break an iterable into chunks of a fixed size. Returns an iterable of iterables.
     Almost-inverse of itertools.chain.from_iterable - passing the output of this into that function will reconstruct the original iterable.
@@ -276,7 +282,7 @@ def chunk(iterable, size):
 def bot_has_permissions(**required):
     """Decorator to check if bot has certain permissions when added to a command"""
 
-    def predicate(ctx):
+    def predicate(ctx: DozerContext):
         """Function to tell the bot if it has the right permissions"""
         given = ctx.channel.permissions_for((ctx.guild or ctx.channel).me)
         missing = [name for name, value in required.items() if getattr(given, name) != value]
@@ -306,11 +312,11 @@ def decorator(func):
 class PrefixHandler:
     """Handles dynamic prefixes"""
 
-    def __init__(self, default_prefix):
+    def __init__(self, default_prefix: str):
         self.default_prefix = default_prefix
-        self.prefix_cache = {}
+        self.prefix_cache: Dict[int, DynamicPrefixEntry] = {}
 
-    def handler(self, bot, message):
+    def handler(self, bot, message: discord.Message):
         """Process the dynamic prefix for each message"""
         dynamic = self.prefix_cache.get(message.guild.id) if message.guild else self.default_prefix
         # <@!> is a nickname mention which discord.py doesn't make by default
@@ -327,7 +333,7 @@ async def refresh(self):
 class DynamicPrefixEntry(db.DatabaseTable):
     """Holds the custom prefixes for guilds"""
     __tablename__ = 'dynamic_prefixes'
-    __uniques__ = 'guild_id'
+    __uniques__ = ['guild_id']
 
     @classmethod
     async def initial_create(cls):
@@ -340,7 +346,7 @@ async def initial_create(cls):
                 PRIMARY KEY (guild_id)
                 )""")
 
-    def __init__(self, guild_id, prefix):
+    def __init__(self, guild_id: int, prefix: str):
         super().__init__()
         self.guild_id = guild_id
         self.prefix = prefix