diff options
| -rw-r--r-- | .travis.yml | 2 | ||||
| -rw-r--r-- | 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 {:<LENGTH} + # to indent it. + start = f"{'':<{len(str(self.ln))+2}}...: " + + if i == len(lines) - 2: + if line.startswith("return"): + line = line[6:].strip() + + # Combine everything + res += (start + line + "\n") + + self.stdout.seek(0) + text = self.stdout.read() + self.stdout.close() + self.stdout = StringIO() + + if text: + res += (text + "\n") + + if out is None: + # No output, return the input statement + return (res, None) + + res += f"Out[{self.ln}]: " + + if isinstance(out, discord.Embed): + # We made an embed? Send that as embed + res += "<Embed>" + res = (res, out) - code = code.strip().strip("\n") - io = StringIO() + else: + if (isinstance(out, str) and + out.startswith("Traceback (most recent call last):\n")): + # Leave out the traceback message + out = "\n" + "\n".join(out.split("\n")[1:]) + + if isinstance(out, str): + pretty = out + else: + pretty = pprint.pformat(out, compact=True, width=60) + + if pretty != str(out): + # We're using the pretty version, start on the next line + res += "\n" + + if pretty.count("\n") > 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): |