From 47136cf9378216c140fae475e188bbee48a822e6 Mon Sep 17 00:00:00 2001 From: martmists Date: Fri, 16 Feb 2018 21:07:49 +0100 Subject: Martmists: Fancier IPython-like eval command (#3) * Fancier IPython-like eval command - Supports Embeds - Pretty-prints - Shortens - Supports catching stdout * Cache pip to speed up travis in the future * Fix snekchek * Update eval.py * Re-order imports * Comply with mentioned issues * Fix travis * Conform to requested edits * Fix travis again * Update eval.py --- .travis.yml | 2 + bot/cogs/eval.py | 212 ++++++++++++++++++++++++++++++++++++++++++------------- 2 files changed, 166 insertions(+), 48 deletions(-) diff --git a/.travis.yml b/.travis.yml index 0782e5a19..ca2161a04 100644 --- a/.travis.yml +++ b/.travis.yml @@ -10,6 +10,8 @@ script: after_success: - python deploy.py +cache: pip + notifications: email: on_success: never diff --git a/bot/cogs/eval.py b/bot/cogs/eval.py index 6fb7f7b06..5a321c2c4 100644 --- a/bot/cogs/eval.py +++ b/bot/cogs/eval.py @@ -1,8 +1,15 @@ # coding=utf-8 +import contextlib +import inspect +import pprint +import re +import textwrap +import traceback from io import StringIO -from discord.ext.commands import AutoShardedBot, Context, command +import discord +from discord.ext.commands import AutoShardedBot, command from bot.constants import OWNER_ROLE from bot.decorators import with_role @@ -12,66 +19,175 @@ from bot.interpreter import Interpreter class EvalCog: # Named this way because a flake8 plugin isn't case-sensitive """ Bot owner only: Evaluate Python code + + Made by martmists """ def __init__(self, bot: AutoShardedBot): self.bot = bot + self.env = {} + self.ln = 0 + self.stdout = StringIO() + self.interpreter = Interpreter(bot) - @command() - @with_role(OWNER_ROLE) - async def eval(self, ctx: Context, *, string: str): - """ - Bot owner only: Evaluate Python code + def _format(self, inp, out): # (str, Any) -> (str, discord.Embed) + self._ = out - Your code may be surrounded in a code fence, but it's not required. - Scope will be preserved - variables set will be present later on. - """ + res = "" - code = string.strip() + # Erase temp input we made + if inp.startswith("_ = "): + inp = inp[4:] + + # Get all non-empty lines + lines = [line for line in inp.split("\n") if line.strip()] + if len(lines) != 1: + lines += [""] + + # Create the input dialog + for i, line in enumerate(lines): + if i == 0: + # Start dialog + start = f"In [{self.ln}]: " - if code.startswith("```") and code.endswith("```"): - if code.startswith("```python"): - code = code[9:-3] - elif code.startswith("```py"): - code = code[5:-3] else: - code = code[3:-3] - elif code.startswith("`") and code.endswith("`"): - code = code[1:-1] + # Indent the 3 dots correctly; + # Normally, it's something like + # In [X]: + # ...: + # + # But if it's + # In [XX]: + # ...: + # + # You can see it doesn't look right. + # This code simply indents the dots + # far enough to align them. + # we first `str()` the line number + # then we get the length + # and do a simple {: 20: + # Text too long, shorten + li = pretty.split("\n") + + pretty = ("\n".join(li[:3]) + # First 3 lines + "\n ...\n" + # Ellipsis to indicate removed lines + "\n".join(li[-3:])) # last 3 lines + + # Add the output + res += pretty + res = (res, None) + + return res # Return (text, embed) + + async def _eval(self, ctx, code): # (discord.Context, str) -> None + self.ln += 1 + + if code.startswith("exit"): + self.ln = 0 + self.env = {} + return await ctx.send("```Reset history!```") + + env = { + "message": ctx.message, + "author": ctx.message.author, + "channel": ctx.channel, + "guild": ctx.guild, + "ctx": ctx, + "self": self, + "bot": self.bot, + "inspect": inspect, + "discord": discord, + "contextlib": contextlib + } + + self.env.update(env) + + # Ignore this shitcode, it works + _code = """ +async def func(): # (None,) -> Any + try: + with contextlib.redirect_stdout(self.stdout): +{0} + if '_' in locals(): + if inspect.isawaitable(_): + _ = await _ + return _ + finally: + self.env.update(locals()) +""".format(textwrap.indent(code, ' ')) try: - rvalue = await self.interpreter.run(code, ctx, io) - except Exception as e: - await ctx.send( - f"{ctx.author.mention} **Code**\n" - f"```py\n{code}```\n\n" - f"**Error**\n```{e}```" - ) - else: - out_message = ( - f"{ctx.author.mention} **Code**\n" - f"```py\n{code}\n```" - ) - - output = io.getvalue() - - if output: - out_message = ( - f"{out_message}\n\n" - f"**Output**\n```{output}```" - ) - - if rvalue is not None: - out_message = ( - f"{out_message}\n\n" - f"**Returned**\n```py\n{repr(rvalue)}\n```" - ) - - await ctx.send(out_message) + exec(_code, self.env) # noqa: B102 pylint: disable=exec-used + func = self.env['func'] + res = await func() + + except Exception: # noqa pylint: disable=broad-except + res = traceback.format_exc() + + out, embed = self._format(code, res) + await ctx.send(f"```py\n{out}```", embed=embed) + + @command() + @with_role(OWNER_ROLE) + async def eval(self, ctx, *, code: str): + """ Run eval in a REPL-like format. """ + code = code.strip("`") + if code.startswith("py\n"): + code = "\n".join(code.split("\n")[1:]) + + if not re.search( # Check if it's an expression + r"^(return|import|for|while|def|class|" + r"from|exit|[a-zA-Z0-9]+\s*=)", code, re.M) and len( + code.split("\n")) == 1: + code = "_ = " + code + + await self._eval(ctx, code) def setup(bot): -- cgit v1.2.3