diff --git a/redbot/core/bot.py b/redbot/core/bot.py index 71f2ff590..8ece19de5 100644 --- a/redbot/core/bot.py +++ b/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) diff --git a/redbot/core/dev_commands.py b/redbot/core/dev_commands.py index 06ae6ceba..2ec5d0b50 100644 --- a/redbot/core/dev_commands.py +++ b/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 = { - "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__", - } - + env = self.get_environment(ctx) 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 = { - "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__", - } - + env = self.get_environment(ctx) 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)