tools.py

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))