Allow users to eval code in either 3.11 or 3.10

To do this we need to track, for each user, the active eval's code, python version and response message. This is because a button press can now also trigger a job to continue. If we did not track these, then editing your own code and then re-evaluating it would trigger the wait_fors in continue_job for each time you pressed the button to change languages.

Co-authored-by: Mark <[email protected]>

1 files changed, 135 insertions, 29 deletions

diff --git a/bot/exts/utils/snekbox.py b/bot/exts/utils/snekbox.py
index 48bb1d3de..6841527df 100644
--- a/bot/exts/utils/snekbox.py
+++ b/bot/exts/utils/snekbox.py
@@ -1,22 +1,23 @@
 import asyncio
 import contextlib
-import datetime
 import re
 from functools import partial
+from operator import attrgetter
 from signal import Signals
 from textwrap import dedent
-from typing import Optional, Tuple
+from typing import Literal, Optional, Tuple
 
 from botcore.utils import scheduling
 from botcore.utils.regex import FORMATTED_CODE_REGEX, RAW_CODE_REGEX
-from discord import AllowedMentions, HTTPException, Message, NotFound, Reaction, User
-from discord.ext.commands import Cog, Command, Context, Converter, command, guild_only
+from discord import AllowedMentions, HTTPException, Interaction, Member, Message, NotFound, Reaction, User, ui
+from discord.ext.commands import Cog, Command, Context, Converter, command, errors, guild_only
 
 from bot.bot import Bot
 from bot.constants import Categories, Channels, Roles, URLs
 from bot.decorators import redirect_output
 from bot.log import get_logger
 from bot.utils import send_to_paste_service
+from bot.utils.lock import LockedResourceError, lock_arg
 from bot.utils.messages import wait_for_deletion
 from bot.utils.services import PasteTooLongError, PasteUploadError
 
@@ -116,6 +117,58 @@ class CodeblockConverter(Converter):
         return codeblocks
 
 
+class PythonVersionSwitcherView(ui.View):
+    """
+    A view to hold the Try in X Python version button.
+
+    A subclass is required to implement a custom interaction_check so only the command invoker
+    can change the Python version used.
+    """
+
+    def __init__(self, owner_id: int) -> None:
+        super().__init__()
+        self.owner_id = owner_id
+
+    async def interaction_check(self, interaction: Interaction) -> bool:
+        """Ensure the user clicking the button is the view owner."""
+        if self.owner_id == interaction.user.id:
+            return True
+
+        await interaction.response.send_message("This is not your button to click!", ephemeral=True)
+        return False
+
+
+class PythonVersionSwitcherButton(ui.Button):
+    """A button that allows users to re-run their eval command in a different Python version."""
+
+    def __init__(
+        self,
+        job_name: str,
+        version_to_switch_to: Literal["3.10", "3.11"],
+        snekbox_cog: "Snekbox",
+        ctx: Context,
+        code: str
+    ) -> None:
+        self.version_to_switch_to = version_to_switch_to
+        super().__init__(label=f"Run in {self.version_to_switch_to}")
+
+        self.snekbox_cog = snekbox_cog
+        self.ctx = ctx
+        self.job_name = job_name
+        self.code = code
+
+    async def callback(self, interaction: Interaction) -> None:
+        """
+        Tell snekbox to re-run the user's code in the alternative Python version.
+
+        Use a task calling snekbox, as run_job is blocking while it waits for edit/reaction on the message.
+        """
+        # Call run_job in a task as blocking here would cause the interaction to be reported as failed by Discord's UI.
+        # Discord.py only ACKs the interaction after this callback finishes.
+        scheduling.create_task(self.snekbox_cog.run_job(self.job_name, self.ctx, self.version_to_switch_to, self.code))
+        await interaction.message.delete()
+
+
 class Snekbox(Cog):
     """Safe evaluation of Python code using Snekbox."""
 
@@ -123,9 +176,37 @@ class Snekbox(Cog):
         self.bot = bot
         self.jobs = {}
 
-    async def post_job(self, code: str, *, args: Optional[list[str]] = None) -> dict:
+    def build_python_version_switcher_view(
+        self,
+        job_name: str,
+        member: Member,
+        current_python_version: Literal["3.10", "3.11"],
+        ctx: Context,
+        code: str
+    ) -> None:
+        """Return a view that allows the user to change what version of Python their code is run on."""
+        if current_python_version == "3.10":
+            alt_python_version = "3.11"
+        else:
+            alt_python_version = "3.10"
+
+        view = PythonVersionSwitcherView(member.id)
+        view.add_item(PythonVersionSwitcherButton(job_name, alt_python_version, self, ctx, code))
+        return view
+
+    async def post_job(
+        self,
+        code: str,
+        python_version: Literal["3.10", "3.11"],
+        *,
+        args: Optional[list[str]] = None
+    ) -> dict:
         """Send a POST request to the Snekbox API to evaluate code and return the results."""
-        url = URLs.snekbox_eval_api
+        if python_version == "3.10":
+            url = URLs.snekbox_eval_api
+        else:
+            url = URLs.snekbox_311_eval_api
+
         data = {"input": code}
 
         if args is not None:
@@ -244,9 +325,11 @@ class Snekbox(Cog):
 
         return output, paste_link
 
+    @lock_arg("snekbox.send_job", "ctx", attrgetter("author.id"), raise_error=True)
     async def send_job(
         self,
         ctx: Context,
+        python_version: Literal["3.10", "3.11"],
         code: str,
         *,
         args: Optional[list[str]] = None,
@@ -258,7 +341,7 @@ class Snekbox(Cog):
         Return the bot response.
         """
         async with ctx.typing():
-            results = await self.post_job(code, args=args)
+            results = await self.post_job(code, python_version, args=args)
             msg, error = self.get_results_message(results, job_name)
 
             if error:
@@ -286,14 +369,15 @@ class Snekbox(Cog):
                 response = await ctx.send("Attempt to circumvent filter detected. Moderator team has been alerted.")
             else:
                 allowed_mentions = AllowedMentions(everyone=False, roles=False, users=[ctx.author])
-                response = await ctx.send(msg, allowed_mentions=allowed_mentions)
+                view = self.build_python_version_switcher_view(job_name, ctx.author, python_version, ctx, code)
+                response = await ctx.send(msg, allowed_mentions=allowed_mentions, view=view)
             scheduling.create_task(wait_for_deletion(response, (ctx.author.id,)), event_loop=self.bot.loop)
 
             log.info(f"{ctx.author}'s {job_name} job had a return code of {results['returncode']}")
         return response
 
     async def continue_job(
-        self, ctx: Context, response: Message, command: Command
+        self, ctx: Context, response: Message, job_name: str
     ) -> tuple[Optional[str], Optional[list[str]]]:
         """
         Check if the job's session should continue.
@@ -318,6 +402,11 @@ class Snekbox(Cog):
                     timeout=10
                 )
 
+                # Ensure the response that's about to be edited is still the most recent.
+                # This could have already been updated via a button press to switch to an alt Python version.
+                if self.jobs[ctx.message.id] != response.id:
+                    return None, None
+
                 code = await self.get_code(new_message, ctx.command)
                 await ctx.message.clear_reaction(REDO_EMOJI)
                 with contextlib.suppress(HTTPException):
@@ -332,7 +421,7 @@ class Snekbox(Cog):
 
             codeblocks = await CodeblockConverter.convert(ctx, code)
 
-            if command is self.timeit_command:
+            if job_name == "timeit":
                 return self.prepare_timeit_input(codeblocks)
             else:
                 return "\n".join(codeblocks), None
@@ -363,18 +452,12 @@ class Snekbox(Cog):
         self,
         job_name: str,
         ctx: Context,
+        python_version: Literal["3.10", "3.11"],
         code: str,
         *,
         args: Optional[list[str]] = None,
     ) -> None:
         """Handles checks, stats and re-evaluation of a snekbox job."""
-        if ctx.author.id in self.jobs:
-            await ctx.send(
-                f"{ctx.author.mention} You've already got a job running - "
-                "please wait for it to finish!"
-            )
-            return
-
         if Roles.helpers in (role.id for role in ctx.author.roles):
             self.bot.stats.incr("snekbox_usages.roles.helpers")
         else:
@@ -390,18 +473,19 @@ class Snekbox(Cog):
         log.info(f"Received code from {ctx.author} for evaluation:\n{code}")
 
         while True:
-            self.jobs[ctx.author.id] = datetime.datetime.now()
-            try:
-                response = await self.send_job(ctx, code, args=args, job_name=job_name)
-            finally:
-                del self.jobs[ctx.author.id]
+            response = await self.send_job(ctx, python_version, code, args=args, job_name=job_name)
 
-            code, args = await self.continue_job(ctx, response, ctx.command)
+            # Store the bot's response message id per invocation, to ensure the `wait_for`s in `continue_job`
+            # don't trigger if the response has already been replaced by a new response.
+            # This can happen when a button is pressed and then original code is edited and re-run.
+            self.jobs[ctx.message.id] = response.id
+
+            code, args = await self.continue_job(ctx, response, job_name)
             if not code:
                 break
             log.info(f"Re-evaluating code from message {ctx.message.id}:\n{code}")
 
-    @command(name="eval", aliases=("e",), usage="<code, ...>")
+    @command(name="eval", aliases=("e",), usage="[python_version] <code, ...>")
     @guild_only()
     @redirect_output(
         destination_channel=Channels.bot_commands,
@@ -410,7 +494,13 @@ class Snekbox(Cog):
         channels=NO_SNEKBOX_CHANNELS,
         ping_user=False
     )
-    async def eval_command(self, ctx: Context, *, code: CodeblockConverter) -> None:
+    async def eval_command(
+        self,
+        ctx: Context,
+        python_version: Optional[Literal["3.10", "3.11"]],
+        *,
+        code: CodeblockConverter
+    ) -> None:
         """
         Run Python code and get the results.
 
@@ -421,12 +511,17 @@ class Snekbox(Cog):
         If multiple codeblocks are in a message, all of them will be joined and evaluated,
         ignoring the text outside of them.
 
+        By default your code is run on Python's 3.11 beta release, to assist with testing. If you
+        run into issues related to this Python version, you can request the bot to use Python
+        3.10 by specifying the `python_version` arg and setting it to `3.10`.
+
         We've done our best to make this sandboxed, but do let us know if you manage to find an
         issue with it!
         """
-        await self.run_job("eval", ctx, "\n".join(code))
+        python_version = python_version or "3.11"
+        await self.run_job("eval", ctx, python_version, "\n".join(code))
 
-    @command(name="timeit", aliases=("ti",), usage="[setup_code] <code, ...>")
+    @command(name="timeit", aliases=("ti",), usage="[python_version] [setup_code] <code, ...>")
     @guild_only()
     @redirect_output(
         destination_channel=Channels.bot_commands,
@@ -435,7 +530,13 @@ class Snekbox(Cog):
         channels=NO_SNEKBOX_CHANNELS,
         ping_user=False
     )
-    async def timeit_command(self, ctx: Context, *, code: CodeblockConverter) -> None:
+    async def timeit_command(
+        self,
+        ctx: Context,
+        python_version: Optional[Literal["3.10", "3.11"]],
+        *,
+        code: CodeblockConverter
+    ) -> None:
         """
         Profile Python Code to find execution time.
 
@@ -446,12 +547,17 @@ class Snekbox(Cog):
         If multiple formatted codeblocks are provided, the first one will be the setup code, which will
         not be timed. The remaining codeblocks will be joined together and timed.
 
+        By default your code is run on Python's 3.11 beta release, to assist with testing. If you
+        run into issues related to this Python version, you can request the bot to use Python
+        3.10 by specifying the `python_version` arg and setting it to `3.10`.
+
         We've done our best to make this sandboxed, but do let us know if you manage to find an
         issue with it!
         """
+        python_version = python_version or "3.11"
         code, args = self.prepare_timeit_input(code)
 
-        await self.run_job("timeit", ctx, code=code, args=args)
+        await self.run_job("timeit", ctx, python_version, code=code, args=args)
 
 
 def predicate_message_edit(ctx: Context, old_msg: Message, new_msg: Message) -> bool: