Skip to content

Commit 0d23bb6

Browse files
committed
Add run_python_code_safely.py as generated by perplexity, plus ruff format, bandit nosec
1 parent 1f728c0 commit 0d23bb6

File tree

1 file changed

+64
-0
lines changed

1 file changed

+64
-0
lines changed
Lines changed: 64 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,64 @@
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

Comments
 (0)