File size: 5,434 Bytes
edf3100 | 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 | import os
import ast
import io
import sys
import numpy as np
import pandas as pd
import scipy
ALLOWED_MODULES = {"numpy", "pandas", "scipy"}
def interpret_python_math_code(python_code: str) -> str:
"""
Interprets a string of Python code to perform math calculations.
Security Note: This function uses exec(). While it attempts to restrict
imports to numpy, pandas, and scipy, and runs with a restricted
global scope, executing arbitrary code always carries risks. Ensure
that input code is from a trusted source or properly sanitized.
The code must only import modules from the allowed list: numpy, pandas, scipy.
Submodules of these (e.g., numpy.linalg, scipy.stats) are permitted.
For example:
'import numpy as np' is allowed.
'from scipy.stats import norm' is allowed.
'import os' is NOT allowed.
To return a result, the code should either:
1. End with an expression (e.g., '1 + 1' or 'np.array([1,2,3]).sum()').
2. Assign the result to a variable named '_result' (e.g., '_result = my_calculation').
Print statements will also be captured and returned along with the result.
"""
# 1. Validate imports using AST
try:
tree = ast.parse(python_code)
for node in tree.body:
if isinstance(node, ast.Import):
for alias in node.names:
root_module = alias.name.split('.')[0]
if root_module not in ALLOWED_MODULES:
return (f"Error: Import of '{alias.name}' is not allowed. "
f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
elif isinstance(node, ast.ImportFrom):
if node.module: # Handles cases like 'from . import something' where module is None
root_module = node.module.split('.')[0]
if root_module not in ALLOWED_MODULES:
return (f"Error: Import from '{node.module}' is not allowed. "
f"Only modules from {list(ALLOWED_MODULES)} are permitted.")
except SyntaxError as e:
return f"Syntax Error in input code: {e}"
# 2. Prepare execution environment
restricted_globals = {
"__builtins__": {
"print": print,
"abs": abs, "round": round, "min": min, "max": max, "sum": sum, "len": len,
"range": range, "zip": zip, "enumerate": enumerate,
"int": int, "float": float, "str": str, "list": list, "dict": dict, "tuple": tuple, "set": set,
"True": True, "False": False, "None": None,
"__import__": __import__, # Add this line
}
# numpy, pandas, scipy are NOT pre-loaded here.
# The user's code `import numpy` will use Python's import mechanism.
# The AST check above is the primary guard.
}
local_vars = {}
# 3. Capture stdout
old_stdout = sys.stdout
redirected_output = io.StringIO()
sys.stdout = redirected_output
# 4. Execute code and retrieve result
calculated_value = None
result_source = ""
output_str = ""
try:
compiled_code = compile(python_code, '<string>', 'exec')
exec(compiled_code, restricted_globals, local_vars)
# Priority 1: Check for '_result' variable
if "_result" in local_vars:
calculated_value = local_vars["_result"]
result_source = "variable '_result'"
# Priority 2: If no _result, and the last AST node was an expression, evaluate it.
elif tree.body and isinstance(tree.body[-1], ast.Expr):
# Ensure the expression node's value is a valid AST object for ast.Expression
if isinstance(tree.body[-1].value, ast.AST):
last_expr_ast = ast.Expression(body=tree.body[-1].value)
# Compile the expression in 'eval' mode
compiled_expr = compile(last_expr_ast, '<string>', 'eval')
# Evaluate in the context of restricted_globals and local_vars (which holds state from exec)
calculated_value = eval(compiled_expr, restricted_globals, local_vars)
result_source = "last expression"
sys.stdout = old_stdout # Restore stdout before getting its value
output_str = redirected_output.getvalue()
if calculated_value is not None:
return f"Result (from {result_source}):\n{calculated_value}\n\nCaptured Output:\n{output_str}".strip()
else:
return f"Executed successfully.\n\nCaptured Output:\n{output_str}\n(No specific result value found via '_result' variable or last expression evaluation.)".strip()
except Exception as e:
if sys.stdout == redirected_output: # Ensure stdout is restored on error too
sys.stdout = old_stdout
output_str = redirected_output.getvalue() # Get any output captured before the error
return f"Execution Error: {type(e).__name__}: {e}\n\nCaptured Output:\n{output_str}".strip()
finally:
# Ensure stdout is always restored
if sys.stdout == redirected_output:
sys.stdout = old_stdout
# Example usage:
if __name__ == "__main__":
code = """
import numpy as np
# import os # This should trigger an error since 'os' is not allowed
arr = np.array([1, 2, 3, 4, 5])
_result = arr.mean()
"""
print(interpret_python_math_code(code))
|