|
| 1 | +import multiprocessing |
| 2 | +import subprocess # nosec B404 |
| 3 | +import sys |
| 4 | +from io import StringIO |
| 5 | + |
| 6 | + |
| 7 | +def run_python_code_safely(python_code, timeout=None): |
| 8 | + """Replacement for subprocess.run that forces 'spawn' context""" |
| 9 | + ctx = multiprocessing.get_context("spawn") |
| 10 | + result_queue = ctx.Queue() |
| 11 | + |
| 12 | + def worker(): |
| 13 | + # Capture stdout/stderr |
| 14 | + old_stdout = sys.stdout |
| 15 | + old_stderr = sys.stderr |
| 16 | + sys.stdout = StringIO() |
| 17 | + sys.stderr = StringIO() |
| 18 | + |
| 19 | + returncode = 0 |
| 20 | + try: |
| 21 | + exec(python_code, {"__name__": "__main__"}) # nosec B102 |
| 22 | + except SystemExit as e: # Handle sys.exit() |
| 23 | + returncode = e.code if isinstance(e.code, int) else 0 |
| 24 | + except Exception: # Capture other exceptions |
| 25 | + import traceback |
| 26 | + |
| 27 | + traceback.print_exc() |
| 28 | + returncode = 1 |
| 29 | + finally: |
| 30 | + # Collect outputs and restore streams |
| 31 | + stdout = sys.stdout.getvalue() |
| 32 | + stderr = sys.stderr.getvalue() |
| 33 | + sys.stdout = old_stdout |
| 34 | + sys.stderr = old_stderr |
| 35 | + result_queue.put((returncode, stdout, stderr)) |
| 36 | + |
| 37 | + process = ctx.Process(target=worker) |
| 38 | + process.start() |
| 39 | + |
| 40 | + try: |
| 41 | + # Wait with timeout support |
| 42 | + process.join(timeout) |
| 43 | + if process.is_alive(): |
| 44 | + process.terminate() |
| 45 | + process.join() |
| 46 | + raise subprocess.TimeoutExpired([sys.executable, "-c", python_code], timeout) |
| 47 | + |
| 48 | + # Get results from queue |
| 49 | + if result_queue.empty(): |
| 50 | + return subprocess.CompletedProcess( |
| 51 | + [sys.executable, "-c", python_code], |
| 52 | + returncode=-999, |
| 53 | + stdout="", |
| 54 | + stderr="Process failed to return results", |
| 55 | + ) |
| 56 | + |
| 57 | + returncode, stdout, stderr = result_queue.get() |
| 58 | + return subprocess.CompletedProcess( |
| 59 | + [sys.executable, "-c", python_code], returncode=returncode, stdout=stdout, stderr=stderr |
| 60 | + ) |
| 61 | + finally: |
| 62 | + # Cleanup if needed |
| 63 | + if process.is_alive(): |
| 64 | + process.kill() |
0 commit comments