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

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):