From 1b777a3689934b2adda8fd75412b4dc6da3abe9d Mon Sep 17 00:00:00 2001 From: Xevion Date: Thu, 18 Feb 2021 01:35:33 -0600 Subject: [PATCH] Split ContestCog into ContestCommandsCog and ContestEventsCog expected_react_deletions and expected_msg_deletions were moved into bot .py There may be more edits included in this that I cannot see. Added docstrings to get_message and fetch_message --- bot/bot.py | 7 +- bot/cogs/contest_commands.py | 168 +++++++++++++++++++ bot/cogs/{contest.py => contest_events.py} | 178 +++------------------ 3 files changed, 192 insertions(+), 161 deletions(-) create mode 100644 bot/cogs/contest_commands.py rename bot/cogs/{contest.py => contest_events.py} (55%) diff --git a/bot/bot.py b/bot/bot.py index 1824c52..27f97c4 100644 --- a/bot/bot.py +++ b/bot/bot.py @@ -1,7 +1,7 @@ import logging from contextlib import contextmanager from datetime import datetime -from typing import ContextManager, List, Optional +from typing import ContextManager, List, Optional, Tuple import discord from discord.ext import commands @@ -23,6 +23,9 @@ class ContestBot(commands.Bot): self.engine = engine self.Session = sessionmaker(bind=engine) + self.expected_msg_deletions: List[int] = [] + self.expected_react_deletions: List[Tuple[int, int]] = [] + @contextmanager def get_session(self, autocommit=True, autoclose=True, rollback=True) -> ContextManager[Session]: """Provides automatic commit and closing of Session with exception rollback.""" @@ -108,10 +111,12 @@ class ContestBot(commands.Bot): await message.add_reaction(self.get_emoji(constants.Emoji.UPVOTE)) def get_message(self, channel_id: int, message_id: int) -> discord.PartialMessage: + """Get a PartialMessage object given raw integer IDs.""" channel: discord.TextChannel = self.get_channel(channel_id) return channel.get_partial_message(message_id) async def fetch_message(self, channel_id: int, message_id: int) -> discord.Message: + """Fetch a full Message object given raw integer IDs.""" channel: discord.TextChannel = self.get_channel(channel_id) return await channel.fetch_message(message_id) diff --git a/bot/cogs/contest_commands.py b/bot/cogs/contest_commands.py new file mode 100644 index 0000000..b6847b3 --- /dev/null +++ b/bot/cogs/contest_commands.py @@ -0,0 +1,168 @@ +import logging + +import discord +from discord.ext import commands +from discord.ext.commands import BucketType, Context, errors + +from bot import checks, constants +from bot.bot import ContestBot +from bot.models import Guild, Period, PeriodStates + +logger = logging.getLogger(__file__) +logger.setLevel(constants.LOGGING_LEVEL) + + +class ContestCommandsCog(commands.Cog): + """ + Manages all commands related to contests. + """ + + def __init__(self, bot: ContestBot) -> None: + self.bot = bot + + @commands.command() + @commands.guild_only() + @checks.privileged() + async def prefix(self, ctx, new_prefix: str): + """Changes the bot's saved prefix.""" + + with self.bot.get_session() as session: + guild: Guild = session.query(Guild).filter_by(id=ctx.guild.id).first() + + if 1 <= len(new_prefix) <= 2: + if guild.prefix == new_prefix: + return await ctx.send(f':no_entry_sign: The prefix is already `{new_prefix}`.') + else: + guild.prefix = new_prefix + return await ctx.send(f':white_check_mark: Prefix changed to `{new_prefix}`.') + else: + return await ctx.send(':no_entry_sign: Invalid argument. Prefix must be 1 or 2 characters long.') + + @commands.command() + @commands.guild_only() + @checks.privileged() + async def submission(self, ctx: Context, new_submission: discord.TextChannel) -> None: + """Changes the bot's saved submission channel.""" + + with self.bot.get_session() as session: + guild: Guild = session.query(Guild).get(ctx.guild.id) + + if guild.submission_channel is not None and guild.submission_channel == new_submission.id: + await ctx.send(f':no_entry_sign: The submission channel is already set to {new_submission.mention}.') + else: + # TODO: Add channel permissions resetting/migration + guild.submission_channel = new_submission.id + await ctx.send(f':white_check_mark: Submission channel changed to {new_submission.mention}.') + + # noinspection PyDunderSlots,PyUnresolvedReferences + @commands.command() + @commands.guild_only() + @commands.has_permissions(send_messages=True, add_reactions=True, read_message_history=True, manage_roles=True) + @commands.max_concurrency(1, per=BucketType.guild, wait=True) + @checks.privileged() + async def advance(self, ctx: Context, duration: float = None, pingback: bool = True) -> None: + """ + Advance the state of the current period pertaining to this Guild. + + :param ctx: + :param duration: If given, the advance command will be repeated once more after the duration (in seconds) has passed. + :param pingback: Whether or not the user should be pinged back when the duration is passed. + """ + # TODO: Ensure that permissions for this command are being correctly tested for. + if duration is not None: assert duration >= 0, "If specified, duration must be more than or equal to zero." + + with self.bot.get_session() as session: + guild: Guild = session.query(Guild).get(ctx.guild.id) + period: Period = guild.current_period + + # Handle non-existent or previously completed period in the current guild + if period is None or not period.active: + if period is None: + overwrite = discord.PermissionOverwrite() + overwrite.send_messages = False + overwrite.add_reactions = False + await self.bot.get_channel(guild.submission_channel).set_permissions(ctx.guild.default_role, overwrite=overwrite) + await ctx.send('Period created, channel permissions set.') + + period = Period(guild_id=guild.id) + session.add(period) + session.commit() + + guild.current_period = period + await ctx.send('New period started - submissions and voting disabled.') + else: + channel: discord.TextChannel = self.bot.get_channel(guild.submission_channel) + target_role: discord.Role = ctx.guild.default_role + # TODO: Research best way to implement contest roles with vagabondit's input + + overwrite = discord.PermissionOverwrite() + overwrite.send_messages = False + overwrite.add_reactions = False + response = 'Permissions unchanged - Period state error.' + + # Handle previous period being completed. + if period.state == PeriodStates.READY: + overwrite.send_messages = True + response = 'Period started, submissions allowed. Advance again to pause.' + # Handle submissions state + elif period.state == PeriodStates.SUBMISSIONS: + response = 'Period paused, submissions disabled. Advance again to start voting.' + # Handle voting state + elif period.state == PeriodStates.PAUSED: + _guild: discord.Guild = ctx.guild + await self.bot.add_voting_reactions(channel=channel, submissions=period.submissions) + overwrite.add_reactions = True + response = 'Period unpaused, reactions allowed. Advance again to stop voting and finalize the tallying.' + # Print period submissions + elif period.state == PeriodStates.VOTING: + response = 'Period stopped. Reactions and submissions disabled. Advance again to start a new period.' + # TODO: Fetch all submissions related to this period and show a embed + + period.advance_state() + + await channel.set_permissions(target_role, overwrite=overwrite) + await ctx.send(response) + + @advance.error + async def advance_error(self, error: errors.CommandError, ctx: Context) -> None: + if isinstance(error, errors.MissingPermissions): + await ctx.send( + 'Check that the bot can actually modify roles, add reactions, see messages and send messages within this channel.') + + # noinspection PyDunderSlots, PyUnresolvedReferences + @commands.command() + @commands.guild_only() + @checks.privileged() + async def close(self, ctx: Context) -> None: + """Closes the current period.""" + with self.bot.get_session() as session: + guild: Guild = session.query(Guild).get(ctx.guild.id) + period: Period = guild.current_period + + if period is None or not period.active: + await ctx.send('No period is currently active.') + else: + overwrite = discord.PermissionOverwrite() + overwrite.send_messages = False + overwrite.add_reactions = False + period.deactivate() + await ctx.send('The current period has been closed.') + + @commands.command() + @commands.guild_only() + async def status(self, ctx: Context) -> None: + """Provides the bot's current state in relation to internal configuration and the server's contest, if active.""" + # TODO: Implement status command + pass + + @commands.command() + @commands.guild_only() + async def leaderboard(self, ctx: Context, count: int = 10, page: int = 0) -> None: + """Prints a leaderboard""" + # TODO: Implement leaderboard command + # TODO: Make interactive and reaction-based + pass + + +def setup(bot) -> None: + bot.add_cog(ContestCommandsCog(bot)) diff --git a/bot/cogs/contest.py b/bot/cogs/contest_events.py similarity index 55% rename from bot/cogs/contest.py rename to bot/cogs/contest_events.py index c2a0946..730bfc0 100644 --- a/bot/cogs/contest.py +++ b/bot/cogs/contest_events.py @@ -3,18 +3,14 @@ from typing import List, Tuple import discord from discord.ext import commands -from discord.ext.commands import BucketType, Context, errors -from bot import checks, constants, helpers +from bot import constants, helpers from bot.bot import ContestBot from bot.models import Guild, Period, PeriodStates, Submission logger = logging.getLogger(__file__) logger.setLevel(constants.LOGGING_LEVEL) -expected_msg_deletions: List[int] = [] -expected_react_deletions: List[Tuple[int, int]] = [] - # TODO: Add command error handling to all commands # TODO: Use embeds in all bot responses @@ -23,154 +19,11 @@ expected_react_deletions: List[Tuple[int, int]] = [] # TODO: Refactor Period into Contest (major) -class ContestCog(commands.Cog): +class ContestEventsCog(commands.Cog): + """Manages all non-command events related to contests.""" def __init__(self, bot: ContestBot): self.bot = bot - @commands.command() - @commands.guild_only() - @checks.privileged() - async def prefix(self, ctx, new_prefix: str): - """Changes the bot's saved prefix.""" - - with self.bot.get_session() as session: - guild: Guild = session.query(Guild).filter_by(id=ctx.guild.id).first() - - if 1 <= len(new_prefix) <= 2: - if guild.prefix == new_prefix: - return await ctx.send(f':no_entry_sign: The prefix is already `{new_prefix}`.') - else: - guild.prefix = new_prefix - return await ctx.send(f':white_check_mark: Prefix changed to `{new_prefix}`.') - else: - return await ctx.send(':no_entry_sign: Invalid argument. Prefix must be 1 or 2 characters long.') - - @commands.command() - @commands.guild_only() - @checks.privileged() - async def submission(self, ctx: Context, new_submission: discord.TextChannel) -> None: - """Changes the bot's saved submission channel.""" - - with self.bot.get_session() as session: - guild: Guild = session.query(Guild).get(ctx.guild.id) - - if guild.submission_channel is not None and guild.submission_channel == new_submission.id: - await ctx.send(f':no_entry_sign: The submission channel is already set to {new_submission.mention}.') - else: - # TODO: Add channel permissions resetting/migration - guild.submission_channel = new_submission.id - await ctx.send(f':white_check_mark: Submission channel changed to {new_submission.mention}.') - - # noinspection PyDunderSlots,PyUnresolvedReferences - @commands.command() - @commands.guild_only() - @commands.has_permissions(send_messages=True, add_reactions=True, read_message_history=True, manage_roles=True) - @commands.cooldown(rate=2, per=5, type=BucketType.guild) - @commands.max_concurrency(1, per=BucketType.guild, wait=False) - @checks.privileged() - async def advance(self, ctx: Context, duration: float = None, pingback: bool = True) -> None: - """ - Advance the state of the current period pertaining to this Guild. - - :param ctx: - :param duration: If given, the advance command will be repeated once more after the duration (in seconds) has passed. - :param pingback: Whether or not the user should be pinged back when the duration is passed. - """ - # TODO: Ensure that permissions for this command are being correctly tested for. - if duration is not None: assert duration >= 0, "If specified, duration must be more than or equal to zero." - - with self.bot.get_session() as session: - guild: Guild = session.query(Guild).get(ctx.guild.id) - period: Period = guild.current_period - - # Handle non-existent or previously completed period in the current guild - if period is None or not period.active: - if period is None: - overwrite = discord.PermissionOverwrite() - overwrite.send_messages = False - overwrite.add_reactions = False - await self.bot.get_channel(guild.submission_channel).set_permissions(ctx.guild.default_role, overwrite=overwrite) - await ctx.send('Period created, channel permissions set.') - - period = Period(guild_id=guild.id) - session.add(period) - session.commit() - - guild.current_period = period - await ctx.send('New period started - submissions and voting disabled.') - else: - channel: discord.TextChannel = self.bot.get_channel(guild.submission_channel) - target_role: discord.Role = ctx.guild.default_role - # TODO: Research best way to implement contest roles with vagabondit's input - - overwrite = discord.PermissionOverwrite() - overwrite.send_messages = False - overwrite.add_reactions = False - response = 'Permissions unchanged - Period state error.' - - # Handle previous period being completed. - if period.state == PeriodStates.READY: - overwrite.send_messages = True - response = 'Period started, submissions allowed. Advance again to pause.' - # Handle submissions state - elif period.state == PeriodStates.SUBMISSIONS: - response = 'Period paused, submissions disabled. Advance again to start voting.' - # Handle voting state - elif period.state == PeriodStates.PAUSED: - _guild: discord.Guild = ctx.guild - await self.bot.add_voting_reactions(channel=channel, submissions=period.submissions) - overwrite.add_reactions = True - response = 'Period unpaused, reactions allowed. Advance again to stop voting and finalize the tallying.' - # Print period submissions - elif period.state == PeriodStates.VOTING: - response = 'Period stopped. Reactions and submissions disabled. Advance again to start a new period.' - # TODO: Fetch all submissions related to this period and show a embed - - period.advance_state() - - await channel.set_permissions(target_role, overwrite=overwrite) - await ctx.send(response) - - @advance.error - async def advance_error(self, error: errors.CommandError, ctx: Context) -> None: - if isinstance(error, errors.MissingPermissions): - await ctx.send( - 'Check that the bot can actually modify roles, add reactions, see messages and send messages within this channel.') - - # noinspection PyDunderSlots, PyUnresolvedReferences - @commands.command() - @commands.guild_only() - @checks.privileged() - async def close(self, ctx: Context) -> None: - """Closes the current period.""" - with self.bot.get_session() as session: - guild: Guild = session.query(Guild).get(ctx.guild.id) - period: Period = guild.current_period - - if period is None or not period.active: - await ctx.send('No period is currently active.') - else: - overwrite = discord.PermissionOverwrite() - overwrite.send_messages = False - overwrite.add_reactions = False - period.deactivate() - await ctx.send('The current period has been closed.') - - @commands.command() - @commands.guild_only() - async def status(self, ctx: Context) -> None: - """Provides the bot's current state in relation to internal configuration and the server's contest, if active.""" - # TODO: Implement status command - pass - - @commands.command() - @commands.guild_only() - async def leaderboard(self, ctx: Context, count: int = 10, page: int = 0) -> None: - """Prints a leaderboard""" - # TODO: Implement leaderboard command - # TODO: Make interactive and reaction-based - pass - @commands.Cog.listener() async def on_message(self, message: discord.Message): if message.author == self.bot.user or message.author.bot or not message.guild: return @@ -184,15 +37,18 @@ class ContestCog(commands.Cog): # Ensure that the submission contains at least one attachment if len(attachments) == 0: - await self.bot.reject(message, f':no_entry_sign: {message.author.mention} Each submission must contain exactly one image.') + await self.bot.reject(message, + f':no_entry_sign: {message.author.mention} Each submission must contain exactly one image.') # Ensure the image contains no more than one attachment elif len(attachments) > 1: - await self.bot.reject(message, f':no_entry_sign: {message.author.mention} Each submission must contain exactly one image.') + await self.bot.reject(message, + f':no_entry_sign: {message.author.mention} Each submission must contain exactly one image.') elif guild.current_period is None: await self.bot.reject(message, f':no_entry_sign: {message.author.mention} A period has not been started. ' f'Submissions should not be allowed at this moment.') elif guild.current_period != PeriodStates.SUBMISSIONS: - logger.warning(f'Valid submission was sent outside of Submissions in {channel.id}/{message.id}. Permissions error? Removing.') + logger.warning( + f'Valid submission was sent outside of Submissions in {channel.id}/{message.id}. Permissions error? Removing.') await message.delete() else: attachment = attachments[0] @@ -202,14 +58,15 @@ class ContestCog(commands.Cog): elif attachment.width is None: await self.bot.reject(message, ':no_entry_sign: Attachment must be a image or video.') else: - last_submission: Submission = session.query(Submission).filter_by(period=guild.current_period, user=message.author.id).first() + last_submission: Submission = session.query(Submission).filter_by(period=guild.current_period, + user=message.author.id).first() if last_submission is not None: # delete last submission submission_msg = await channel.fetch_message(last_submission.id) if submission_msg is None: logger.error(f'Unexpected: submission message {last_submission.id} could not be found.') else: - expected_msg_deletions.append(submission_msg.id) + self.bot.expected_msg_deletions.append(submission_msg.id) await submission_msg.delete() logger.info(f'Old submission deleted. {last_submission.id} (Old) -> {message.id} (New)') @@ -227,8 +84,8 @@ class ContestCog(commands.Cog): await self.bot.wait_until_ready() # Ignore messages we delete - if payload.message_id in expected_msg_deletions: - expected_msg_deletions.remove(payload.message_id) + if payload.message_id in self.bot.expected_msg_deletions: + self.bot.expected_msg_deletions.remove(payload.message_id) return with self.bot.get_session() as session: @@ -269,6 +126,7 @@ class ContestCog(commands.Cog): if payload.channel_id == guild.submission_channel: channel: discord.TextChannel = self.bot.get_channel(payload.channel_id) message: discord.PartialMessage = channel.get_partial_message(payload.message_id) + if helpers.is_upvote(payload.emoji): submission: Submission = session.query(Submission).get(payload.message_id) if submission is None: @@ -291,8 +149,8 @@ class ContestCog(commands.Cog): """Deal with reactions we remove or removed manually by users.""" # Skip reactions we removed ourselves. try: - index = expected_react_deletions.index((payload.message_id, payload.emoji.id)) - del expected_react_deletions[index] + index = self.bot.expected_react_deletions.index((payload.message_id, payload.emoji.id)) + del self.bot.expected_react_deletions[index] logger.debug(f'Skipping expected reaction removal {payload.message_id}.') return except ValueError: @@ -340,4 +198,4 @@ class ContestCog(commands.Cog): def setup(bot) -> None: - bot.add_cog(ContestCog(bot)) + bot.add_cog(ContestEventsCog(bot))