[Dev] Customizable environment values (#4667)

* Make the dev env flexible

* Fix rst format in docstrings

* Reproduce current behaviour for _ in repl

* Prevent adding existing or reserved names

* Fix typo with environment

* Docstring changes

Apply suggestions from code review

Co-authored-by: jack1142 <6032823+jack1142@users.noreply.github.com>

* Get env before loop

* Hey I'm not the only one doing typos

* Keep new messages in env

* Clear exception of stack frames

Co-authored-by: jack1142 <6032823+jack1142@users.noreply.github.com>

* Include the `channel` variable in the reserved names

* And we're also missing `discord` :)

Co-authored-by: jack1142 <6032823+jack1142@users.noreply.github.com>

redbot/core/bot.py

@@ -282,6 +282,91 @@ class RedBase(
         """
         self._help_formatter = commands.help.RedHelpFormatter()
 
+    def add_dev_env_value(self, name: str, value: Callable[[commands.Context], Any]):
+        """
+        Add a custom variable to the dev environment (``[p]debug``, ``[p]eval``, and ``[p]repl`` commands).
+        If dev mode is disabled, nothing will happen.
+
+        .. admonition:: Example
+
+            .. code-block:: python
+
+                class MyCog(commands.Cog):
+                    def __init__(self, bot):
+                        self.bot = bot
+                        bot.add_dev_env_value("mycog", lambda ctx: self)
+                        bot.add_dev_env_value("mycogdata", lambda ctx: self.settings[ctx.guild.id])
+
+                    def cog_unload(self):
+                        self.bot.remove_dev_env_value("mycog")
+                        self.bot.remove_dev_env_value("mycogdata")
+
+            Once your cog is loaded, the custom variables ``mycog`` and ``mycogdata``
+            will be included in the environment of dev commands.
+
+        Parameters
+        ----------
+        name: str
+            The name of your custom variable.
+        value: Callable[[commands.Context], Any]
+            The function returning the value of the variable.
+            It must take a `commands.Context` as its sole parameter
+
+        Raise
+        -----
+        TypeError
+            ``value`` argument isn't a callable.
+        ValueError
+            The passed callable takes no or more than one argument.
+        RuntimeError
+            The name of the custom variable is either reserved by a variable
+            from the default environment or already taken by some other custom variable.
+        """
+        signature = inspect.signature(value)
+        if len(signature.parameters) != 1:
+            raise ValueError("Callable must take exactly one argument for context")
+        dev = self.get_cog("Dev")
+        if dev is None:
+            return
+        if name in [
+            "bot",
+            "ctx",
+            "channel",
+            "author",
+            "guild",
+            "message",
+            "asyncio",
+            "aiohttp",
+            "discord",
+            "commands",
+            "_",
+            "__name__",
+            "__builtins__",
+        ]:
+            raise RuntimeError(f"The name {name} is reserved for default environement.")
+        if name in dev.env_extensions:
+            raise RuntimeError(f"The name {name} is already used.")
+        dev.env_extensions[name] = value
+
+    def remove_dev_env_value(self, name: str):
+        """
+        Remove a custom variable from the dev environment.
+
+        Parameters
+        ----------
+        name: str
+            The name of the custom variable.
+
+        Raise
+        -----
+        KeyError
+            The custom variable was never set.
+        """
+        dev = self.get_cog("Dev")
+        if dev is None:
+            return
+        del dev.env_extensions[name]
+
     def get_command(self, name: str) -> Optional[commands.Command]:
         com = super().get_command(name)
         assert com is None or isinstance(com, commands.Command)

redbot/core/dev_commands.py

@@ -45,6 +45,7 @@ class Dev(commands.Cog):
         super().__init__()
         self._last_result = None
         self.sessions = {}
+        self.env_extensions = {}
 
     @staticmethod
     def async_compile(source, filename, mode):
@@ -92,6 +93,29 @@ class Dev(commands.Cog):
         token = ctx.bot.http.token
         return re.sub(re.escape(token), "[EXPUNGED]", input_, re.I)
 
+    def get_environment(self, ctx: commands.Context) -> dict:
+        env = {
+            "bot": ctx.bot,
+            "ctx": ctx,
+            "channel": ctx.channel,
+            "author": ctx.author,
+            "guild": ctx.guild,
+            "message": ctx.message,
+            "asyncio": asyncio,
+            "aiohttp": aiohttp,
+            "discord": discord,
+            "commands": commands,
+            "_": self._last_result,
+            "__name__": "__main__",
+        }
+        for name, value in self.env_extensions.items():
+            try:
+                env[name] = value(ctx)
+            except Exception as e:
+                traceback.clear_frames(e.__traceback__)
+                env[name] = e
+        return env
+
     @commands.command()
     @checks.is_owner()
     async def debug(self, ctx, *, code):
@@ -115,21 +139,7 @@ class Dev(commands.Cog):
             commands - redbot.core.commands
             _        - The result of the last dev command.
         """
-        env = {
+        env = self.get_environment(ctx)
-            "bot": ctx.bot,
-            "ctx": ctx,
-            "channel": ctx.channel,
-            "author": ctx.author,
-            "guild": ctx.guild,
-            "message": ctx.message,
-            "asyncio": asyncio,
-            "aiohttp": aiohttp,
-            "discord": discord,
-            "commands": commands,
-            "_": self._last_result,
-            "__name__": "__main__",
-        }
-
         code = self.cleanup_code(code)
 
         try:
@@ -169,21 +179,7 @@ class Dev(commands.Cog):
             commands - redbot.core.commands
             _        - The result of the last dev command.
         """
-        env = {
+        env = self.get_environment(ctx)
-            "bot": ctx.bot,
-            "ctx": ctx,
-            "channel": ctx.channel,
-            "author": ctx.author,
-            "guild": ctx.guild,
-            "message": ctx.message,
-            "asyncio": asyncio,
-            "aiohttp": aiohttp,
-            "discord": discord,
-            "commands": commands,
-            "_": self._last_result,
-            "__name__": "__main__",
-        }
-
         body = self.cleanup_code(body)
         stdout = io.StringIO()
 
@@ -224,19 +220,6 @@ class Dev(commands.Cog):
         backtick. This includes codeblocks, and as such multiple lines can be
         evaluated.
         """
-        variables = {
-            "ctx": ctx,
-            "bot": ctx.bot,
-            "message": ctx.message,
-            "guild": ctx.guild,
-            "channel": ctx.channel,
-            "author": ctx.author,
-            "asyncio": asyncio,
-            "_": None,
-            "__builtins__": __builtins__,
-            "__name__": "__main__",
-        }
-
         if ctx.channel.id in self.sessions:
             if self.sessions[ctx.channel.id]:
                 await ctx.send(
@@ -250,6 +233,9 @@ class Dev(commands.Cog):
                 )
             return
 
+        env = self.get_environment(ctx)
+        env["__builtins__"] = __builtins__
+        env["_"] = None
         self.sessions[ctx.channel.id] = True
         await ctx.send(
             _(
@@ -287,8 +273,7 @@ class Dev(commands.Cog):
                     await ctx.send(self.get_syntax_error(e))
                     continue
 
-            variables["message"] = response
+            env["message"] = response
-
             stdout = io.StringIO()
 
             msg = ""
@@ -296,9 +281,9 @@ class Dev(commands.Cog):
             try:
                 with redirect_stdout(stdout):
                     if executor is None:
-                        result = types.FunctionType(code, variables)()
+                        result = types.FunctionType(code, env)()
                     else:
-                        result = executor(code, variables)
+                        result = executor(code, env)
                     result = await self.maybe_await(result)
             except:
                 value = stdout.getvalue()
@@ -307,7 +292,7 @@ class Dev(commands.Cog):
                 value = stdout.getvalue()
                 if result is not None:
                     msg = "{}{}".format(value, result)
-                    variables["_"] = result
+                    env["_"] = result
                 elif value:
                     msg = "{}".format(value)