From f18b0a93557ba4f85f8aaacf2e5966e0842ab0ee Mon Sep 17 00:00:00 2001 From: Walter Kuppens Date: Sat, 25 May 2019 14:58:41 -0700 Subject: [PATCH] Switched to using decorators for commands --- cho_client.py | 64 ++++++++++++++++++++++++++++++++++-- cho_commands.py | 86 +++++-------------------------------------------- cho_game.py | 8 +++-- cho_utils.py | 35 ++++++++++++++++++++ 4 files changed, 109 insertions(+), 84 deletions(-) diff --git a/cho_client.py b/cho_client.py index ae491b4..b562f55 100644 --- a/cho_client.py +++ b/cho_client.py @@ -108,6 +108,64 @@ async def on_message(self, message: Message): async def on_error(self, event_name, *args, **kwargs): """Logs exceptions to the bot's log.""" - LOGGER.error( - "Received uncaught exception:\n\n%s" - % traceback.format_exc()) + stack_trace = traceback.format_exc() + LOGGER.error("Received uncaught exception:\n\n%s", stack_trace) + + async def handle_command(self, message): + """Called when a Cho command is received from a user. + + :param m message: + :type m: discord.message.Message + """ + + guild_id = message.guild.id + + # This is a good opportunity to make sure the guild we're getting a + # command from is setup properly in the database. + guild_query_results = sql.guild.get_guild(self.engine, guild_id) + if not guild_query_results: + LOGGER.info("Got command from new guild: %s", guild_id) + sql.guild.create_guild(self.engine, guild_id) + config = {} + else: + _, config = guild_query_results + + # TODO: Come up with a better way to split up arguments. If we want to + # support flags in the future this might need to be done using a real + # argument parser. + args = message.content.split() + + # Handle cho invocations with no command. + if len(args) < 2: + await message.channel.send( + "You didn't specify a command. If you want to " + "start a game use the \"start\" command." + ) + return + + command = args[1].lower() + + # Process commands that are marked for global usage. + for global_command, func in cho_utils.GLOBAL_COMMANDS.items(): + if global_command == command: + await func(self, message, args, config) + return + + # Anything not handled above must be done in the configured channel. + if not cho_utils.is_message_from_trivia_channel(message, config): + await message.channel.send( + "Sorry, I can't be summoned into this channel. Please go " + "to the trivia channel for this server." + ) + return + + # Process commands that are marked for channel-only usage. + for channel_command, func in cho_utils.CHANNEL_COMMANDS.items(): + if channel_command == command: + await func(self, message, args, config) + return + + await message.channel.send( + "I'm afraid I don't know that command. If you want to " + "start a game use the \"start\" command." + ) diff --git a/cho_commands.py b/cho_commands.py index ff8080f..6847bbc 100644 --- a/cho_commands.py +++ b/cho_commands.py @@ -19,9 +19,10 @@ import logging import re import discord -import cho_utils import sql.guild +from cho_utils import cho_command + CMD_START = "start" CMD_STOP = "stop" CMD_SCOREBOARD = "scoreboard" @@ -38,71 +39,7 @@ class ChoCommandsMixin(): """Contains command handler functions for ChoClient.""" - async def handle_command(self, message): - """Called when a Cho command is received from a user. - - :param m message: - :type m: discord.message.Message - """ - - guild_id = message.guild.id - - # This is a good opportunity to make sure the guild we're getting a - # command from is setup properly in the database. - guild_query_results = sql.guild.get_guild(self.engine, guild_id) - if not guild_query_results: - LOGGER.info("Got command from new guild: %s", guild_id) - sql.guild.create_guild(self.engine, guild_id) - config = {} - else: - _, config = guild_query_results - - args = message.content.split() - if len(args) < 2: - await message.channel.send( - "You didn't specify a command. If you want to " - "start a game use the \"start\" command." - ) - return - - command = args[1].lower() - - if command == CMD_HELP: - await self._handle_help(message, args, config) - return - - # Admin commands should be processed anywhere. - if command == CMD_SET_CHANNEL: - await self._handle_set_channel(message, args, config) - return - elif command == CMD_SET_PREFIX: - await self._handle_set_prefix(message, args, config) - return - - # Anything not handled above must be done in the configured channel. - if not cho_utils.is_message_from_trivia_channel(message, config): - await message.channel.send( - "Sorry, I can't be summoned into this channel. Please go " - "to the trivia channel for this server." - ) - return - - # Trivia channel-only commands. - if command == CMD_START: - await self._handle_start_command(message, args, config) - return - elif command == CMD_STOP: - await self._handle_stop_command(message, args, config) - return - elif command == CMD_SCOREBOARD: - await self._handle_scoreboard_command(message, args, config) - return - - await message.channel.send( - "I'm afraid I don't know that command. If you want to " - "start a game use the \"start\" command." - ) - + @cho_command(CMD_HELP) async def _handle_help(self, message, args, config): """Responds with help to teach users about the bot's functions. @@ -154,6 +91,7 @@ async def _handle_help(self, message, args, config): ) await message.channel.send(embed=embed) + @cho_command(CMD_START, kind="channel") async def _handle_start_command(self, message, args, config): """Starts a new game at the request of a user. @@ -179,6 +117,7 @@ async def _handle_start_command(self, message, args, config): ) await self._start_game(message.guild, message.channel) + @cho_command(CMD_STOP, kind="channel") async def _handle_stop_command(self, message, args, config): """Stops the current game at the request of the user. @@ -207,6 +146,7 @@ async def _handle_stop_command(self, message, args, config): "one first." ) + @cho_command(CMD_SCOREBOARD, kind="channel") async def _handle_scoreboard_command(self, message, args, config): """Displays a scoreboard at the request of the user. @@ -249,6 +189,7 @@ async def _handle_scoreboard_command(self, message, args, config): "Currently no scores are available. Try playing a game to " "get some scores in the scoreboard.") + @cho_command(CMD_SET_CHANNEL, admin_only=True) async def _handle_set_channel(self, message, args, config): """Updates the trivia channel configuration for the guild. @@ -264,12 +205,6 @@ async def _handle_set_channel(self, message, args, config): ) return - if not cho_utils.is_admin(message.author, message.channel): - await message.channel.send( - "Sorry, only administrators can move me." - ) - return - guild_id = message.guild.id trivia_channel_id = args[2] trivia_channel_re_match = DISCORD_CHANNEL_REGEX.match( @@ -289,6 +224,7 @@ async def _handle_set_channel(self, message, args, config): "The trivia channel is now in {}.".format(trivia_channel_id) ) + @cho_command(CMD_SET_PREFIX, admin_only=True) async def _handle_set_prefix(self, message, args, config): """Updates the prefix used for the guild. @@ -304,12 +240,6 @@ async def _handle_set_prefix(self, message, args, config): ) return - if not cho_utils.is_admin(message.author, message.channel): - await message.channel.send( - "Sorry, only administrators can change the prefix." - ) - return - guild_id = message.guild.id new_prefix = args[2] diff --git a/cho_game.py b/cho_game.py index 614cb09..59f100f 100644 --- a/cho_game.py +++ b/cho_game.py @@ -18,13 +18,15 @@ import asyncio import logging -import cho_utils -import sql.guild -import sql.scoreboard from discord.channel import TextChannel from discord.guild import Guild from discord.message import Message + +import cho_utils +import sql.guild +import sql.scoreboard + from game_state import GameState SHORT_WAIT_SECS = 5 diff --git a/cho_utils.py b/cho_utils.py index a311011..e170184 100644 --- a/cho_utils.py +++ b/cho_utils.py @@ -17,6 +17,9 @@ """Core code that controls Cho's behavior.""" import logging + +from collections import OrderedDict + import jellyfish from discord.channel import TextChannel @@ -27,6 +30,38 @@ LOGGER = logging.getLogger("cho") +GLOBAL_COMMANDS = OrderedDict() +CHANNEL_COMMANDS = OrderedDict() + + +def cho_command(command, kind="global", admin_only=False): + """Marks a function as a runnable command.""" + + def decorator(func): + def wrapper(*args, **kwargs): + if admin_only: + message = args[1] + + if not is_admin(message.author, message.channel): + return message.channel.send( + "Sorry, only administrators run that command." + ) + + return func(*args, **kwargs) + + return func(*args, **kwargs) + + if kind == "global": + GLOBAL_COMMANDS[command] = wrapper + elif kind == "channel": + CHANNEL_COMMANDS[command] = wrapper + else: + raise ValueError("Unknown cho command type passed in decorator.") + + return wrapper + + return decorator + def get_prefix(config: dict = None) -> str: """Gets the prefix for the specified guild.