|
| 1 | +import multiprocessing |
| 2 | +import subprocess # nosec B404 |
| 3 | +import sys |
| 4 | +import traceback |
| 5 | +from io import StringIO |
| 6 | + |
| 7 | + |
| 8 | +class Worker: |
| 9 | + def __init__(self, python_code, result_queue): |
| 10 | + self.python_code = python_code |
| 11 | + self.result_queue = result_queue |
| 12 | + |
| 13 | + def __call__(self): |
| 14 | + # Capture stdout/stderr |
| 15 | + old_stdout = sys.stdout |
| 16 | + old_stderr = sys.stderr |
| 17 | + sys.stdout = StringIO() |
| 18 | + sys.stderr = StringIO() |
| 19 | + |
| 20 | + try: |
| 21 | + exec(self.python_code, {"__name__": "__main__"}) # nosec B102 |
| 22 | + returncode = 0 |
| 23 | + except SystemExit as e: # Handle sys.exit() |
| 24 | + returncode = e.code if isinstance(e.code, int) else 0 |
| 25 | + except BaseException: |
| 26 | + traceback.print_exc() |
| 27 | + returncode = 1 |
| 28 | + finally: |
| 29 | + # Collect outputs and restore streams |
| 30 | + stdout = sys.stdout.getvalue() |
| 31 | + stderr = sys.stderr.getvalue() |
| 32 | + sys.stdout = old_stdout |
| 33 | + sys.stderr = old_stderr |
| 34 | + try: # noqa: SIM105 |
| 35 | + self.result_queue.put((returncode, stdout, stderr)) |
| 36 | + except Exception: # nosec B110 |
| 37 | + # If the queue is broken (e.g., parent gone), best effort logging |
| 38 | + pass |
| 39 | + |
| 40 | + |
| 41 | +def run_python_code_safely(python_code, *, timeout=None): |
| 42 | + """Run Python code in a spawned subprocess, capturing stdout/stderr/output.""" |
| 43 | + ctx = multiprocessing.get_context("spawn") |
| 44 | + result_queue = ctx.SimpleQueue() |
| 45 | + process = ctx.Process(target=Worker(python_code, result_queue)) |
| 46 | + process.start() |
| 47 | + |
| 48 | + try: |
| 49 | + process.join(timeout) |
| 50 | + if process.is_alive(): |
| 51 | + process.terminate() |
| 52 | + process.join() |
| 53 | + return subprocess.CompletedProcess( |
| 54 | + args=[sys.executable, "-c", python_code], |
| 55 | + returncode=-9, |
| 56 | + stdout="", |
| 57 | + stderr=f"Process timed out after {timeout} seconds and was terminated.", |
| 58 | + ) |
| 59 | + |
| 60 | + if result_queue.empty(): |
| 61 | + return subprocess.CompletedProcess( |
| 62 | + args=[sys.executable, "-c", python_code], |
| 63 | + returncode=-999, |
| 64 | + stdout="", |
| 65 | + stderr="Process exited without returning results.", |
| 66 | + ) |
| 67 | + |
| 68 | + returncode, stdout, stderr = result_queue.get() |
| 69 | + return subprocess.CompletedProcess( |
| 70 | + args=[sys.executable, "-c", python_code], |
| 71 | + returncode=returncode, |
| 72 | + stdout=stdout, |
| 73 | + stderr=stderr, |
| 74 | + ) |
| 75 | + |
| 76 | + finally: |
| 77 | + try: |
| 78 | + result_queue.close() |
| 79 | + result_queue.join_thread() |
| 80 | + except Exception: # nosec B110 |
| 81 | + pass |
| 82 | + if process.is_alive(): |
| 83 | + process.kill() |
| 84 | + process.join() |
0 commit comments