File size: 4,524 Bytes
ec9f5b2 |
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 122 123 124 125 126 127 128 |
import datetime
import io
import time
from contextlib import redirect_stdout, redirect_stderr
import uuid
from matplotlib import pyplot as plt
import pandas as pd
import numpy as np
import traceback
import seaborn as sns
plt.style.use('seaborn-v0_8-whitegrid')
class PythonDataAnalysisExecutor:
"""
Simplified Python code execution environment for data analysis
"""
def __init__(self, timeout_seconds=30):
self.timeout = timeout_seconds
self.safe_globals = {
'__builtins__': {
'print': print, 'range': range, 'len': len,
'str': str, 'int': int, 'float': float, 'bool': bool,
'list': list, 'dict': dict, 'tuple': tuple, 'set': set,
'min': min, 'max': max, 'sum': sum, 'abs': abs,
'round': round, 'zip': zip, 'enumerate': enumerate,
'__import__': __import__
},
'pd': pd, 'np': np,
'matplotlib.pyplot': plt,
'seaborn': sns,
'uuid': uuid.uuid4,
'datetime': datetime, 'time': time,
'DataFrame': pd.DataFrame, 'Series': pd.Series
}
self.last_result = None
def _validate_code(self, code):
"""Basic security checks"""
forbidden = ['sys.', 'open(', 'exec(', 'eval(']
if any(f in code for f in forbidden):
raise ValueError("Potentially unsafe code detected")
def execute(self, code: str) -> dict:
"""
Execute Python code safely with timeout
Returns dict with: success, output, error, execution_time, variables, result
"""
result = {
'success': False,
'output': '',
'error': '',
'execution_time': 0,
'variables': {},
'result': None # This will capture the last expression result
}
start_time = time.time()
output = io.StringIO()
try:
self._validate_code(code)
with redirect_stdout(output), redirect_stderr(output):
# Split into lines and handle last expression
lines = [line for line in code.split('\n') if line.strip()]
if lines:
# Execute all but last line normally
if len(lines) > 1:
exec('\n'.join(lines[:-1]), self.safe_globals)
# Handle last line specially to capture its value
last_line = lines[-1].strip()
if last_line:
# If it's an expression (not assignment or control structure)
if not (last_line.startswith((' ', '\t')) or
last_line.split()[0] in ('if', 'for', 'while', 'def', 'class') or
'=' in last_line):
self.last_result = eval(last_line, self.safe_globals)
result['result'] = self.last_result
output.write(str(self.last_result) + '\n')
else:
exec(last_line, self.safe_globals)
result['output'] = output.getvalue()
result['variables'] = {
k: v for k, v in self.safe_globals.items()
if not k.startswith('__') and k in code
}
result['success'] = True
except Exception as e:
result['error'] = f"{str(e)}\n{traceback.format_exc()}"
result['execution_time'] = time.time() - start_time
return result
def run_analysis(code: str, timeout=20) -> dict:
"""Simplified interface for code execution"""
executor = PythonDataAnalysisExecutor(timeout_seconds=timeout)
return executor.execute(code)
# Example usage
# if __name__ == "__main__":
# analysis_code = """
# import datetime
# print(datetime.datetime.now())
# """
# result = run_analysis(analysis_code)
# # Improved output formatting
# if result['success']:
# print("Execution successful")
# print("Execution time:", result['execution_time'], "seconds")
# print("Output:", result['output'].strip())
# print("Result:", result['result'])
# print("Variables:", list(result['variables'].keys()))
# else:
# print("Execution failed")
# print("Error:", result['error'])
|