Switched to using decorators for commands

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."
+        )

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]
 

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

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.