File size: 5,533 Bytes
3ede55d | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 | #!/usr/bin/env python3
import os, sys, subprocess, tempfile, json, traceback, time, resource
class ExecutionSandbox:
def __init__(self, engine=None, timeout=5, max_memory_mb=256):
self.engine = engine
self.timeout = timeout
self.max_memory = max_memory_mb * 1024 * 1024
self.execution_history = []
def _run_in_subprocess(self, code, test_input=None):
with tempfile.NamedTemporaryFile(mode='w', suffix='.py', delete=False) as f:
f.write(code)
tmp = f.name
try:
script = f"""
import sys, json, traceback, builtins
import resource
resource.setrlimit(resource.RLIMIT_CPU, ({self.timeout}, {self.timeout}))
resource.setrlimit(resource.RLIMIT_AS, ({self.max_memory}, {self.max_memory}))
_original = builtins.__import__
def _restricted(name, *args, **kwargs):
blocked = ['os.system', 'subprocess', 'socket', 'urllib', 'http.client', 'ftplib']
if any(b in name for b in blocked):
raise ImportError(f"Blocked: {{name}}")
return _original(name, *args, **kwargs)
builtins.__import__ = _restricted
result = {{"stdout": "", "stderr": "", "output": None, "error": None, "success": False}}
old_out, old_err = sys.stdout, sys.stderr
try:
from io import StringIO
sys.stdout, sys.stderr = StringIO(), StringIO()
ns = {{"__builtins__": builtins}}
with open("{tmp}") as f:
exec(f.read(), ns)
for name, obj in ns.items():
if callable(obj) and name in ["main", "run", "solve", "test"]:
try:
result["output"] = obj("{test_input or ''}") if "{test_input or ''}" else obj()
except Exception as e:
result["error"] = str(e)
result["stdout"] = sys.stdout.getvalue()
result["stderr"] = sys.stderr.getvalue()
result["success"] = result["error"] is None
except Exception as e:
result["error"] = traceback.format_exc()
result["success"] = False
finally:
sys.stdout, sys.stderr = old_out, old_err
print(json.dumps(result))
"""
proc = subprocess.run([sys.executable, '-c', script], capture_output=True, text=True, timeout=self.timeout+2, cwd='/tmp')
out = proc.stdout.strip()
if out:
try:
return json.loads(out.split('\n')[-1])
except:
return {"stdout": proc.stdout, "stderr": proc.stderr, "error": proc.stderr or "Failed", "success": False}
return {"stdout": "", "stderr": proc.stderr, "error": proc.stderr or "No output", "success": False}
except subprocess.TimeoutExpired:
return {"error": f"Timeout after {self.timeout}s", "success": False}
except Exception as e:
return {"error": str(e), "success": False}
finally:
try: os.unlink(tmp)
except: pass
def execute_and_verify(self, code, expected_output=None, test_input=None):
result = self._run_in_subprocess(code, test_input)
if expected_output is not None:
actual = result.get("output")
result["verified"] = True
result["match"] = actual == expected_output
result["expected"] = expected_output
result["actual"] = actual
else:
result["verified"] = False
if self.engine and hasattr(self.engine, '_update_pheromone'):
success = result.get("success") and result.get("match", True)
for idx in range(1000):
self.engine._update_pheromone(idx, success)
self.execution_history.append({"success": result.get("success"), "match": result.get("match")})
return result
def auto_fix(self, code, error, max_attempts=3):
if not self.engine:
return code, {"fixed": False}
current = code
for attempt in range(max_attempts):
prompt = f"Fix this Python code. Error: {error}\n\n```python\n{current}\n```\n\nProvide ONLY the corrected code."
try:
fixed = self.engine.generate(prompt, max_tokens=800)
if "```python" in fixed:
fixed = fixed.split("```python")[1].split("```")[0]
elif "```" in fixed:
fixed = fixed.split("```")[1].split("```")[0]
result = self._run_in_subprocess(fixed)
if result.get("success"):
return fixed, {"fixed": True, "attempts": attempt + 1}
current = fixed
error = result.get("error", "Unknown")
except Exception as e:
return current, {"fixed": False, "reason": str(e)}
return current, {"fixed": False, "reason": f"Failed after {max_attempts}"}
def get_stats(self):
total = len(self.execution_history)
if total == 0:
return {"total": 0, "success_rate": 0}
successes = sum(1 for h in self.execution_history if h.get("success"))
matches = sum(1 for h in self.execution_history if h.get("match"))
return {"total": total, "successes": successes, "success_rate": successes/total*100, "matches": matches}
if __name__ == "__main__":
s = ExecutionSandbox()
r1 = s.execute_and_verify("def add(a,b): return a+b\nprint(add(2,3))")
print(f"Test 1: success={r1['success']}, stdout={r1.get('stdout','')[:50]}")
r2 = s.execute_and_verify("def div(a,b): return a/b\ndiv(10,0)")
print(f"Test 2: success={r2['success']}, error={r2.get('error','')[:50]}")
print(f"Stats: {s.get_stats()}")
|