diff --git a/cogs/code.py b/cogs/code.py new file mode 100644 index 0000000..c9205f6 --- /dev/null +++ b/cogs/code.py @@ -0,0 +1,78 @@ +import discord +from discord.ext import commands +from pyston import File, PystonClient +from pyston.exceptions import * + + +class Code(commands.Cog): + def __init__(self, bot): + self.bot = bot + + @discord.message_command(description="Execute code") + async def execute(self, ctx, message: discord.Message): + await ctx.defer() # Allow response to take longer than 3 seconds + + code_blocks = get_code_blocks_with_languages(message.content) + + responses = [] + for block in code_blocks: + try: + output = await execute_code(block["language"], block["code"]) + result = f"```{str(output)}```" + except InvalidLanguage: + result = "Invalid language" + except TooManyRequests: + result = "Too many requests. Please slow down" + except (ExecutionError, InternalServerError, UnexpectedError): + result = "An unknown error has occured ¯\_(ツ)_/¯" + + responses.append(f'```{block["language"]}\n{block["code"]}```{result}') + + if len(responses) == 0: + await ctx.respond( + "This command must be used on a message containing a code block with a language specified" + ) + + response = "\n".join(responses) + await ctx.respond(response) + + +async def execute_code(language, code): + client = PystonClient() + + language = convert_language_alias(language) + + # Check if language is valid + await client.get_runtimes(language) # Raises InvalidLanguage if language is invalid + + # Execute code + result = await client.execute( + language, [File(code)] + ) # Raises TooManyRequests, ExecutionError, InternalServerError or UnexpectedError + return result + + +# Get code blocks and languages from a message string +def get_code_blocks_with_languages(txt): + result = [] + blocks = txt.split("```") + i = 1 + while i < len(blocks) - 1: + lang_and_code = blocks[i].split("\n", 1) + result.append({"language": lang_and_code[0], "code": lang_and_code[1]}) + i += 2 + + return result + + +# Convert language to a format that works in Piston +def convert_language_alias(lang): + if lang.lower() == "c#": + return "csharp" + if lang.lower() == "cplusplus": + return "c++" + return lang.lower() + + +def setup(bot): + bot.add_cog(Code(bot)) diff --git a/requirements.txt b/requirements.txt index eff0fe0..d3f8674 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,4 +1,5 @@ aiohttp==3.8.1 py-cord==2.0.0b1 psutil==5.8.0 -requests==2.28.1 \ No newline at end of file +requests==2.28.1 +aiopyston==1.2.1