Spaces:
Sleeping
Sleeping
File size: 2,410 Bytes
0fe6e4d | 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 | import pandas as pd
import numpy as np
import traceback
import ast
#####validate generated code
def validate_code(code):
tree = ast.parse(code)
forbidden_calls = {
"eval",
"exec",
"open",
"__import__",
"compile",
"input"
}
forbidden_modules = {
"os",
"sys",
"subprocess",
"shutil",
"socket",
"pathlib"
}
for node in ast.walk(tree):
# import x
if isinstance(node, ast.Import):
for alias in node.names:
if alias.name not in ["pandas", "numpy"]:
raise ValueError(
f"Import not allowed: {alias.name}"
)
# from x import y
if isinstance(node, ast.ImportFrom):
if node.module not in ["pandas", "numpy"]:
raise ValueError(
f"Import not allowed: {node.module}"
)
# function calls
if isinstance(node, ast.Call):
if isinstance(node.func, ast.Name):
if node.func.id in forbidden_calls:
raise ValueError(
f"Forbidden function: {node.func.id}"
)
# attribute access
if isinstance(node, ast.Attribute):
if node.attr.startswith("__"):
raise ValueError(
"Dunder access forbidden"
)
# direct module usage
if isinstance(node, ast.Name):
if node.id in forbidden_modules:
raise ValueError(
f"Forbidden module: {node.id}"
)
return True
def execute_generated_code(code, df):
try:
validate_code(code)
local_vars = {
"df": df,
"pd": pd,
"np": np
}
#exec(code, {}, local_vars)
safe_globals = {
"__builtins__": {}
}
exec(
code,
safe_globals,
local_vars
)
return {
"success": True,
"result": local_vars.get("result")
}
except Exception as e:
return {
"success": False,
"error": str(e)
} |