Update codeblock regex

The snekbox implementation of the codeblock regex
was incorporated. This now correctly parses the
`code` and ``code`` markdown discord allows.

You can also use multiple code blocks with text interrupting it and
it will process the different code blocks as one continuous code block.

1 files changed, 36 insertions, 1 deletions

diff --git a/bot/exts/internal_eval/_internal_eval.py b/bot/exts/internal_eval/_internal_eval.py
index f7a0946b..45bfbdc3 100644
--- a/bot/exts/internal_eval/_internal_eval.py
+++ b/bot/exts/internal_eval/_internal_eval.py
@@ -17,6 +17,23 @@ log = logging.getLogger("rattlesnake.exts.admin_tools.internal_eval")
 
 CODEBLOCK_REGEX = re.compile(r"(^```(py(thon)?)?\n)|(```$)")
 
+FORMATTED_CODE_REGEX = re.compile(
+    r"(?P<delim>(?P<block>```)|``?)"        # code delimiter: 1-3 backticks; (?P=block) only matches if it's a block
+    r"(?(block)(?:(?P<lang>[a-z]+)\n)?)"    # if we're in a block, match optional language (only letters plus newline)
+    r"(?:[ \t]*\n)*"                        # any blank (empty or tabs/spaces only) lines before the code
+    r"(?P<code>.*?)"                        # extract all code inside the markup
+    r"\s*"                                  # any more whitespace before the end of the code markup
+    r"(?P=delim)",                          # match the exact same delimiter from the start again
+    re.DOTALL | re.IGNORECASE               # "." also matches newlines, case insensitive
+)
+
+RAW_CODE_REGEX = re.compile(
+    r"^(?:[ \t]*\n)*"                       # any blank (empty or tabs/spaces only) lines before the code
+    r"(?P<code>.*?)"                        # extract all the rest as code
+    r"\s*$",                                # any trailing whitespace until the end of the string
+    re.DOTALL                               # "." also matches newlines
+)
+
 
 class InternalEval(commands.Cog):
     """Top secret code evaluation for admins and owners."""
@@ -141,7 +158,25 @@ class InternalEval(commands.Cog):
     @with_role(Roles.admin)
     async def eval(self, ctx: commands.Context, *, code: str) -> None:
         """Run eval in a REPL-like format."""
-        code = CODEBLOCK_REGEX.sub("", code.strip())
+
+        if match := list(FORMATTED_CODE_REGEX.finditer(code)):
+            blocks = [block for block in match if block.group("block")]
+
+            if len(blocks) > 1:
+                code = '\n'.join(block.group("code") for block in blocks)
+                info = "several code blocks"
+            else:
+                match = match[0] if len(blocks) == 0 else blocks[0]
+                code, block, lang, delim = match.group("code", "block", "lang", "delim")
+                if block:
+                    info = (f"'{lang}' highlighted" if lang else "plain") + " code block"
+                else:
+                    info = f"{delim}-enclosed inline code"
+        else:
+            code = RAW_CODE_REGEX.fullmatch(code).group("code")
+            info = "unformatted or badly formatted code"
+
+        code = textwrap.dedent(code)
         await self._eval(ctx, code)
 
     @internal_group.command(name='reset', aliases=("clear", "exit", "r", "c"))