Snekbox: fix re-eval when '!eval' is removed from edited message

The previous parsing method was naïve in assuming there would always
be something preceding the code (e.g. the '!eval' command invocation)
delimited by a space. Now it will only split if it's sure the eval
command was used in the edited message.

1 files changed, 21 insertions, 1 deletions

diff --git a/bot/cogs/snekbox.py b/bot/cogs/snekbox.py
index cff7c5786..454836921 100644
--- a/bot/cogs/snekbox.py
+++ b/bot/cogs/snekbox.py
@@ -232,7 +232,7 @@ class Snekbox(Cog):
                     timeout=10
                 )
 
-                code = new_message.content.split(' ', maxsplit=1)[1]
+                code = await self.get_code(new_message)
                 await ctx.message.clear_reactions()
                 with contextlib.suppress(HTTPException):
                     await response.delete()
@@ -243,6 +243,26 @@ class Snekbox(Cog):
 
             return code
 
+    async def get_code(self, message: Message) -> Optional[str]:
+        """
+        Return the code from `message` to be evaluated.
+
+        If the message is an invocation of the eval command, return the first argument or None if it
+        doesn't exist. Otherwise, return the full content of the message.
+        """
+        log.trace(f"Getting context for message {message.id}.")
+        new_ctx = await self.bot.get_context(message)
+
+        if new_ctx.command is self.eval_command:
+            log.trace(f"Message {message.id} invokes eval command.")
+            split = message.content.split(maxsplit=1)
+            code = split[1] if len(split) > 1 else None
+        else:
+            log.trace(f"Message {message.id} does not invoke eval command.")
+            code = message.content
+
+        return code
+
     @command(name="eval", aliases=("e",))
     @guild_only()
     @in_channel(Channels.bot_commands, hidden_channels=(Channels.esoteric,), bypass_roles=EVAL_ROLES)