diff options
| -rw-r--r-- | bot/cogs/snekbox.py | 22 |
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) |