executor.py

import abc
import sys
import io
import ast
import subprocess
from contextlib import redirect_stdout
from loguru import logger

# Configure logger
logger.remove()
logger.add(
    sys.stderr,
    format="<green>{time}</green> <level>{message}</level>",
    colorize=True,
)

class Executor(abc.ABC):
    @abc.abstractmethod
    def execute(self, code: str) -> str:
        pass


class PythonExecutor(Executor):
    locals = {}

    def execute(self, code: str) -> str:
        logger.info("Executing Python code:\n{}", code)
        output = io.StringIO()

        # Parse the code into an AST.
        tree = ast.parse(code, mode="exec")

        try:
            # Redirect standard output to our StringIO instance.
            with redirect_stdout(output):
                for node in tree.body:
                    # Compile and execute each node.
                    exec(
                        compile(
                            ast.Module(body=[node], type_ignores=[]), "<ast>", "exec"
                        ),
                        None,
                        PythonExecutor.locals,
                    )

                    # If the node is an expression, print its result.
                    if isinstance(node, ast.Expr):
                        eval_result = eval(
                            compile(ast.Expression(body=node.value), "<ast>", "eval"),
                            None,
                            PythonExecutor.locals,
                        )
                        if eval_result is not None:
                            print(eval_result)
        except Exception as e:
            logger.error("Error executing Python code: {}", str(e))
            return f"Error: {str(e)}"

        # Retrieve the output and return it.
        return output.getvalue()


class CppExecutor(Executor):
    def execute(self, code: str) -> str:
        with open("script.cpp", "w") as f:
            f.write(code)
        try:
            logger.info("Compiling and running C++ code.")
            subprocess.run(["g++", "script.cpp", "-o", "script"], check=True)
            output = subprocess.run(
                ["./script"], capture_output=True, text=True, check=True
            )
            return output.stdout
        except subprocess.CalledProcessError as e:
            logger.error("Error executing C++ code: {}", e.stderr)
            return f"Compilation/Execution Error:\n{e.stderr}"


class RustExecutor(Executor):
    def execute(self, code: str) -> str:
        with open("script.rs", "w") as f:
            f.write(code)
        try:
            logger.info("Compiling and running Rust code.")
            subprocess.run(["rustc", "script.rs", "-o", "script"], check=True)
            output = subprocess.run(
                ["./script"], capture_output=True, text=True, check=True
            )
            return output.stdout
        except subprocess.CalledProcessError as e:
            logger.error("Error executing Rust code: {}", e.stderr)
            return f"Compilation/Execution Error:\n{e.stderr}"