rp/pr/tools/python_exec.py

import contextlib
import os
import traceback
from io import StringIO


def python_exec(code, python_globals, cwd=None):
    """Execute Python code and capture the output.

    Args:
        code: The Python code to execute.
        python_globals: Dictionary of global variables for execution.
        cwd: Working directory for execution.

    Returns:
        Dict with status and output, or error information.
    """
    try:
        original_cwd = None
        if cwd:
            original_cwd = os.getcwd()
            os.chdir(cwd)

        output = StringIO()
        with contextlib.redirect_stdout(output):
            exec(code, python_globals)

        if original_cwd:
            os.chdir(original_cwd)

        return {"status": "success", "output": output.getvalue()}
    except Exception as e:
        if original_cwd:
            os.chdir(original_cwd)
        return {"status": "error", "error": str(e), "traceback": traceback.format_exc()}