quantalogic 0.35.0__py3-none-any.whl → 0.50.0__py3-none-any.whl

quantalogic/utils/python_interpreter.py

@@ -1,28 +1,37 @@
 import ast
 import builtins
 import textwrap
-from typing import Any, List, Dict, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple
+
 
 # Exception used to signal a "return" from a function call.
 class ReturnException(Exception):
     def __init__(self, value: Any) -> None:
         self.value: Any = value
 
+
 # Exceptions used for loop control.
 class BreakException(Exception):
     pass
 
+
 class ContinueException(Exception):
     pass
 
+
 # The main interpreter class.
 class ASTInterpreter:
     def __init__(
-        self,
-        allowed_modules: List[str],
-        env_stack: Optional[List[Dict[str, Any]]] = None,
-        source: Optional[str] = None
+        self, allowed_modules: List[str], env_stack: Optional[List[Dict[str, Any]]] = None, source: Optional[str] = None
     ) -> None:
+        """
+        Initialize the AST interpreter with restricted module access and environment stack.
+
+        Args:
+            allowed_modules: List of module names allowed to be imported.
+            env_stack: Optional pre-existing environment stack; if None, a new one is created.
+            source: Optional source code string for error reporting.
+        """
         self.allowed_modules: List[str] = allowed_modules
         self.modules: Dict[str, Any] = {}
         # Import only the allowed modules.
@@ -61,7 +70,7 @@ class ASTInterpreter:
         else:
             self.source_lines = None
 
-        # NEW: Add standard Decimal features if allowed.
+        # Add standard Decimal features if allowed.
         if "decimal" in self.modules:
             dec = self.modules["decimal"]
             self.env_stack[0]["Decimal"] = dec.Decimal
@@ -77,8 +86,9 @@ class ASTInterpreter:
         globals: Optional[Dict[str, Any]] = None,
         locals: Optional[Dict[str, Any]] = None,
         fromlist: Tuple[str, ...] = (),
-        level: int = 0
+        level: int = 0,
     ) -> Any:
+        """Restrict imports to only allowed modules."""
         if name not in self.allowed_modules:
             error_msg = f"Import Error: Module '{name}' is not allowed. Only {self.allowed_modules} are permitted."
             raise ImportError(error_msg)
@@ -86,14 +96,14 @@ class ASTInterpreter:
 
     # Helper: create a new interpreter instance using a given environment stack.
     def spawn_from_env(self, env_stack: List[Dict[str, Any]]) -> "ASTInterpreter":
+        """Spawn a new interpreter with the provided environment stack."""
         return ASTInterpreter(
-            self.allowed_modules,
-            env_stack,
-            source="\n".join(self.source_lines) if self.source_lines else None
+            self.allowed_modules, env_stack, source="\n".join(self.source_lines) if self.source_lines else None
         )
 
     # Look up a variable in the chain of environment frames.
     def get_variable(self, name: str) -> Any:
+        """Retrieve a variable's value from the environment stack."""
         for frame in reversed(self.env_stack):
             if name in frame:
                 return frame[name]
@@ -101,10 +111,12 @@ class ASTInterpreter:
 
     # Always assign to the most local environment.
     def set_variable(self, name: str, value: Any) -> None:
+        """Set a variable in the most local environment frame."""
         self.env_stack[-1][name] = value
 
     # Used for assignment targets. This handles names and destructuring.
     def assign(self, target: ast.AST, value: Any) -> None:
+        """Assign a value to a target (name, tuple, attribute, or subscript)."""
         if isinstance(target, ast.Name):
             # If current frame declares the name as global, update global frame.
             if "__global_names__" in self.env_stack[-1] and target.id in self.env_stack[-1]["__global_names__"]:
@@ -127,15 +139,15 @@ class ASTInterpreter:
             else:
                 total = len(value)
                 before = target.elts[:star_index]
-                after = target.elts[star_index+1:]
+                after = target.elts[star_index + 1 :]
                 if len(before) + len(after) > total:
                     raise ValueError("Unpacking mismatch")
                 for i, elt2 in enumerate(before):
                     self.assign(elt2, value[i])
                 starred_count = total - len(before) - len(after)
-                self.assign(target.elts[star_index].value, value[len(before):len(before)+starred_count])
+                self.assign(target.elts[star_index].value, value[len(before) : len(before) + starred_count])
                 for j, elt2 in enumerate(after):
-                    self.assign(elt2, value[len(before)+starred_count+j])
+                    self.assign(elt2, value[len(before) + starred_count + j])
         elif isinstance(target, ast.Attribute):
             obj = self.visit(target.value)
             setattr(obj, target.attr, value)
@@ -148,6 +160,7 @@ class ASTInterpreter:
 
     # Main visitor dispatch.
     def visit(self, node: ast.AST) -> Any:
+        """Dispatch to the appropriate visitor method for the AST node."""
         method_name: str = "visit_" + node.__class__.__name__
         method = getattr(self, method_name, self.generic_visit)
         try:
@@ -162,12 +175,11 @@ class ASTInterpreter:
             context_line = ""
             if self.source_lines and 1 <= lineno <= len(self.source_lines):
                 context_line = self.source_lines[lineno - 1]
-            raise Exception(
-                f"Error line {lineno}, col {col}:\n{context_line}\nDescription: {str(e)}"
-            ) from e
+            raise Exception(f"Error line {lineno}, col {col}:\n{context_line}\nDescription: {str(e)}") from e
 
     # Fallback for unsupported nodes.
     def generic_visit(self, node: ast.AST) -> Any:
+        """Handle unsupported AST nodes with an error."""
         lineno = getattr(node, "lineno", None)
         context_line = ""
         if self.source_lines and lineno is not None and 1 <= lineno <= len(self.source_lines):
@@ -186,14 +198,19 @@ class ASTInterpreter:
             module_name: str = alias.name
             asname: str = alias.asname if alias.asname is not None else module_name
             if module_name not in self.allowed_modules:
-                raise Exception(f"Import Error: Module '{module_name}' is not allowed. Only {self.allowed_modules} are permitted.")
+                raise Exception(
+                    f"Import Error: Module '{module_name}' is not allowed. Only {self.allowed_modules} are permitted."
+                )
             self.set_variable(asname, self.modules[module_name])
 
     def visit_ImportFrom(self, node: ast.ImportFrom) -> None:
+        """Process 'from ... import ...' statements with restricted module access."""
         if not node.module:
             raise Exception("Import Error: Missing module name in 'from ... import ...' statement")
         if node.module not in self.allowed_modules:
-            raise Exception(f"Import Error: Module '{node.module}' is not allowed. Only {self.allowed_modules} are permitted.")
+            raise Exception(
+                f"Import Error: Module '{node.module}' is not allowed. Only {self.allowed_modules} are permitted."
+            )
         for alias in node.names:
             if alias.name == "*":
                 raise Exception("Import Error: 'from ... import *' is not supported.")
@@ -205,8 +222,7 @@ class ASTInterpreter:
     def visit_ListComp(self, node: ast.ListComp) -> List[Any]:
         """
         Process a list comprehension, e.g., [elt for ... in ... if ...].
-        The comprehension is executed in a new local frame that inherits the
-        current environment.
+        The comprehension is executed in a new local frame that inherits the current environment.
         """
         result: List[Any] = []
         # Copy the current top-level frame for the comprehension scope.
@@ -234,19 +250,22 @@ class ASTInterpreter:
 
     # --- Other node visitors below ---
     def visit_Module(self, node: ast.Module) -> Any:
-        # Execute all statements then return the 'result' variable if set.
+        """Execute module body and return 'result' or last value."""
         last_value: Any = None
         for stmt in node.body:
             last_value = self.visit(stmt)
         return self.env_stack[0].get("result", last_value)
 
     def visit_Expr(self, node: ast.Expr) -> Any:
+        """Evaluate an expression statement."""
         return self.visit(node.value)
 
     def visit_Constant(self, node: ast.Constant) -> Any:
+        """Return the value of a constant node."""
         return node.value
 
     def visit_Name(self, node: ast.Name) -> Any:
+        """Handle variable name lookups or stores."""
         if isinstance(node.ctx, ast.Load):
             return self.get_variable(node.id)
         elif isinstance(node.ctx, ast.Store):
@@ -255,6 +274,7 @@ class ASTInterpreter:
             raise Exception("Unsupported context for Name")
 
     def visit_BinOp(self, node: ast.BinOp) -> Any:
+        """Evaluate binary operations."""
         left: Any = self.visit(node.left)
         right: Any = self.visit(node.right)
         op = node.op
@@ -271,7 +291,7 @@ class ASTInterpreter:
         elif isinstance(op, ast.Mod):
             return left % right
         elif isinstance(op, ast.Pow):
-            return left ** right
+            return left**right
         elif isinstance(op, ast.LShift):
             return left << right
         elif isinstance(op, ast.RShift):
@@ -286,6 +306,7 @@ class ASTInterpreter:
             raise Exception("Unsupported binary operator: " + str(op))
 
     def visit_UnaryOp(self, node: ast.UnaryOp) -> Any:
+        """Evaluate unary operations."""
         operand: Any = self.visit(node.operand)
         op = node.op
         if isinstance(op, ast.UAdd):
@@ -300,11 +321,13 @@ class ASTInterpreter:
             raise Exception("Unsupported unary operator: " + str(op))
 
     def visit_Assign(self, node: ast.Assign) -> None:
+        """Handle assignment statements."""
         value: Any = self.visit(node.value)
         for target in node.targets:
             self.assign(target, value)
 
     def visit_AugAssign(self, node: ast.AugAssign) -> Any:
+        """Handle augmented assignments (e.g., +=)."""
         # If target is a Name, get its current value from the environment.
         if isinstance(node.target, ast.Name):
             current_val: Any = self.get_variable(node.target.id)
@@ -325,7 +348,7 @@ class ASTInterpreter:
         elif isinstance(op, ast.Mod):
             result = current_val % right_val
         elif isinstance(op, ast.Pow):
-            result = current_val ** right_val
+            result = current_val**right_val
         elif isinstance(op, ast.BitAnd):
             result = current_val & right_val
         elif isinstance(op, ast.BitOr):
@@ -342,6 +365,7 @@ class ASTInterpreter:
         return result
 
     def visit_Compare(self, node: ast.Compare) -> bool:
+        """Evaluate comparison operations."""
         left: Any = self.visit(node.left)
         for op, comparator in zip(node.ops, node.comparators):
             right: Any = self.visit(comparator)
@@ -364,13 +388,13 @@ class ASTInterpreter:
                 if not (left >= right):
                     return False
             elif isinstance(op, ast.Is):
-                if not (left is right):
+                if left is not right:
                     return False
             elif isinstance(op, ast.IsNot):
                 if not (left is not right):
                     return False
             elif isinstance(op, ast.In):
-                if not (left in right):
+                if left not in right:
                     return False
             elif isinstance(op, ast.NotIn):
                 if not (left not in right):
@@ -381,6 +405,7 @@ class ASTInterpreter:
         return True
 
     def visit_BoolOp(self, node: ast.BoolOp) -> bool:
+        """Evaluate boolean operations (and/or)."""
         if isinstance(node.op, ast.And):
             for value in node.values:
                 if not self.visit(value):
@@ -395,6 +420,7 @@ class ASTInterpreter:
             raise Exception("Unsupported boolean operator: " + str(node.op))
 
     def visit_If(self, node: ast.If) -> Any:
+        """Handle if statements."""
         if self.visit(node.test):
             branch = node.body
         else:
@@ -409,6 +435,7 @@ class ASTInterpreter:
         return result
 
     def visit_While(self, node: ast.While) -> None:
+        """Handle while loops."""
         while self.visit(node.test):
             try:
                 for stmt in node.body:
@@ -421,6 +448,7 @@ class ASTInterpreter:
             self.visit(stmt)
 
     def visit_For(self, node: ast.For) -> None:
+        """Handle for loops."""
         iter_obj: Any = self.visit(node.iter)
         for item in iter_obj:
             self.assign(node.target, item)
@@ -435,53 +463,66 @@ class ASTInterpreter:
             self.visit(stmt)
 
     def visit_Break(self, node: ast.Break) -> None:
+        """Handle break statements."""
         raise BreakException()
 
     def visit_Continue(self, node: ast.Continue) -> None:
+        """Handle continue statements."""
         raise ContinueException()
 
     def visit_FunctionDef(self, node: ast.FunctionDef) -> None:
+        """Define a function and store it in the environment."""
         # Capture the current env_stack for a closure without copying inner dicts.
-        closure: List[Dict[str, Any]] = self.env_stack[:]  # <-- changed here
+        closure: List[Dict[str, Any]] = self.env_stack[:]
         func = Function(node, closure, self)
         self.set_variable(node.name, func)
 
     def visit_Call(self, node: ast.Call) -> Any:
+        """Handle function calls."""
         func = self.visit(node.func)
         args: List[Any] = [self.visit(arg) for arg in node.args]
         kwargs: Dict[str, Any] = {kw.arg: self.visit(kw.value) for kw in node.keywords}
         return func(*args, **kwargs)
 
     def visit_Return(self, node: ast.Return) -> None:
+        """Handle return statements."""
         value: Any = self.visit(node.value) if node.value is not None else None
         raise ReturnException(value)
 
     def visit_Lambda(self, node: ast.Lambda) -> Any:
-        closure: List[Dict[str, Any]] = self.env_stack[:]  # <-- changed here
+        """Define a lambda function."""
+        closure: List[Dict[str, Any]] = self.env_stack[:]
         return LambdaFunction(node, closure, self)
 
     def visit_List(self, node: ast.List) -> List[Any]:
+        """Evaluate list literals."""
         return [self.visit(elt) for elt in node.elts]
 
     def visit_Tuple(self, node: ast.Tuple) -> Tuple[Any, ...]:
+        """Evaluate tuple literals."""
         return tuple(self.visit(elt) for elt in node.elts)
 
     def visit_Dict(self, node: ast.Dict) -> Dict[Any, Any]:
+        """Evaluate dictionary literals."""
         return {self.visit(k): self.visit(v) for k, v in zip(node.keys, node.values)}
 
     def visit_Set(self, node: ast.Set) -> set:
+        """Evaluate set literals."""
         return set(self.visit(elt) for elt in node.elts)
 
     def visit_Attribute(self, node: ast.Attribute) -> Any:
+        """Handle attribute access."""
         value: Any = self.visit(node.value)
         return getattr(value, node.attr)
 
     def visit_Subscript(self, node: ast.Subscript) -> Any:
+        """Handle subscript operations (e.g., list indexing)."""
         value: Any = self.visit(node.value)
         slice_val: Any = self.visit(node.slice)
         return value[slice_val]
 
     def visit_Slice(self, node: ast.Slice) -> slice:
+        """Evaluate slice objects."""
         lower: Any = self.visit(node.lower) if node.lower else None
         upper: Any = self.visit(node.upper) if node.upper else None
         step: Any = self.visit(node.step) if node.step else None
@@ -489,20 +530,23 @@ class ASTInterpreter:
 
     # For compatibility with older AST versions.
     def visit_Index(self, node: ast.Index) -> Any:
+        """Handle index nodes for older AST compatibility."""
         return self.visit(node.value)
 
     # Visitor for Pass nodes.
     def visit_Pass(self, node: ast.Pass) -> None:
-        # Simply ignore 'pass' statements.
+        """Handle pass statements (do nothing)."""
         return None
 
     def visit_TypeIgnore(self, node: ast.TypeIgnore) -> None:
+        """Handle type ignore statements (do nothing)."""
         pass
 
     def visit_Try(self, node: ast.Try) -> Any:
+        """Handle try-except blocks."""
         result: Any = None
         exc_info: Optional[tuple] = None
-        
+
         try:
             for stmt in node.body:
                 result = self.visit(stmt)
@@ -541,23 +585,24 @@ class ASTInterpreter:
                     if exc_info:
                         raise exc_info[1]
                     raise
-        
+
         return result
 
     def visit_Nonlocal(self, node: ast.Nonlocal) -> None:
+        """Handle nonlocal statements (minimal support)."""
         # Minimal support – assume these names exist in an outer frame.
         return None
 
     def visit_JoinedStr(self, node: ast.JoinedStr) -> str:
-        # Support f-string: concatenate all parts.
+        """Handle f-strings by concatenating all parts."""
         return "".join(self.visit(value) for value in node.values)
 
     def visit_FormattedValue(self, node: ast.FormattedValue) -> str:
-        # Format the embedded expression.
+        """Handle formatted values within f-strings."""
         return str(self.visit(node.value))
 
     def visit_GeneratorExp(self, node: ast.GeneratorExp) -> Any:
-        # Process a generator expression.
+        """Handle generator expressions."""
         def generator():
             base_frame: Dict[str, Any] = self.env_stack[-1].copy()
             self.env_stack.append(base_frame)
@@ -575,28 +620,29 @@ class ASTInterpreter:
                         if all(self.visit(if_clause) for if_clause in comp.ifs):
                             yield from rec(gen_idx + 1)
                         self.env_stack.pop()
+
             gen = list(rec(0))
             self.env_stack.pop()
             for val in gen:
                 yield val
+
         return generator()
 
     def visit_ClassDef(self, node: ast.ClassDef):
+        """Handle class definitions."""
         base_frame = self.env_stack[-1].copy()
         self.env_stack.append(base_frame)
         try:
             for stmt in node.body:
                 self.visit(stmt)
-            class_dict = {
-                k: v for k, v in self.env_stack[-1].items() 
-                if k not in ["__builtins__"]
-            }
+            class_dict = {k: v for k, v in self.env_stack[-1].items() if k not in ["__builtins__"]}
         finally:
             self.env_stack.pop()
         new_class = type(node.name, (), class_dict)
         self.set_variable(node.name, new_class)
 
     def visit_With(self, node: ast.With):
+        """Handle with statements."""
         for item in node.items:
             ctx = self.visit(item.context_expr)
             val = ctx.__enter__()
@@ -612,21 +658,26 @@ class ASTInterpreter:
                 ctx.__exit__(None, None, None)
 
     def visit_Raise(self, node: ast.Raise):
+        """Handle raise statements."""
         exc = self.visit(node.exc) if node.exc else None
         if exc:
             raise exc
         raise Exception("Raise with no exception specified")
 
     def visit_Global(self, node: ast.Global):
+        """Handle global statements."""
         self.env_stack[-1].setdefault("__global_names__", set()).update(node.names)
 
     def visit_IfExp(self, node: ast.IfExp):
+        """Handle ternary if expressions."""
         return self.visit(node.body) if self.visit(node.test) else self.visit(node.orelse)
 
     def visit_DictComp(self, node: ast.DictComp):
+        """Handle dictionary comprehensions."""
         result = {}
         base_frame = self.env_stack[-1].copy()
         self.env_stack.append(base_frame)
+
         def rec(gen_idx: int):
             if gen_idx == len(node.generators):
                 key = self.visit(node.key)
@@ -641,14 +692,17 @@ class ASTInterpreter:
                     if all(self.visit(if_clause) for if_clause in comp.ifs):
                         rec(gen_idx + 1)
                     self.env_stack.pop()
+
         rec(0)
         self.env_stack.pop()
         return result
 
     def visit_SetComp(self, node: ast.SetComp):
+        """Handle set comprehensions."""
         result = set()
         base_frame = self.env_stack[-1].copy()
         self.env_stack.append(base_frame)
+
         def rec(gen_idx: int):
             if gen_idx == len(node.generators):
                 result.add(self.visit(node.elt))
@@ -661,13 +715,16 @@ class ASTInterpreter:
                     if all(self.visit(if_clause) for if_clause in comp.ifs):
                         rec(gen_idx + 1)
                     self.env_stack.pop()
+
         rec(0)
         self.env_stack.pop()
         return result
 
+
 # Class to represent a user-defined function.
 class Function:
     def __init__(self, node: ast.FunctionDef, closure: List[Dict[str, Any]], interpreter: ASTInterpreter) -> None:
+        """Initialize a user-defined function."""
         self.node: ast.FunctionDef = node
         # Shallow copy to support recursion.
         self.closure: List[Dict[str, Any]] = self.env_stack_reference(closure)
@@ -675,9 +732,11 @@ class Function:
 
     # Helper to simply return the given environment stack (shallow copy of list refs).
     def env_stack_reference(self, env_stack: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Return a shallow copy of the environment stack."""
         return env_stack[:]  # shallow
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the function with given arguments."""
         new_env_stack: List[Dict[str, Any]] = self.closure[:]
         local_frame: Dict[str, Any] = {}
         # Bind the function into its own local frame for recursion.
@@ -703,21 +762,26 @@ class Function:
 
     # Add __get__ to support method binding.
     def __get__(self, instance: Any, owner: Any):
+        """Support method binding for instance methods."""
         def method(*args: Any, **kwargs: Any) -> Any:
             return self(instance, *args, **kwargs)
         return method
 
+
 # Class to represent a lambda function.
 class LambdaFunction:
     def __init__(self, node: ast.Lambda, closure: List[Dict[str, Any]], interpreter: ASTInterpreter) -> None:
+        """Initialize a lambda function."""
         self.node: ast.Lambda = node
         self.closure: List[Dict[str, Any]] = self.env_stack_reference(closure)
         self.interpreter: ASTInterpreter = interpreter
 
     def env_stack_reference(self, env_stack: List[Dict[str, Any]]) -> List[Dict[str, Any]]:
+        """Return a shallow copy of the environment stack."""
         return env_stack[:]
 
     def __call__(self, *args: Any, **kwargs: Any) -> Any:
+        """Execute the lambda function with given arguments."""
         new_env_stack: List[Dict[str, Any]] = self.closure[:]
         local_frame: Dict[str, Any] = {}
         if len(args) < len(self.node.args.args):
@@ -732,9 +796,12 @@ class LambdaFunction:
         new_interp: ASTInterpreter = self.interpreter.spawn_from_env(new_env_stack)
         return new_interp.visit(self.node.body)
 
+
 # The main function to interpret an AST.
 def interpret_ast(ast_tree: Any, allowed_modules: list[str], source: str = "") -> Any:
+    """Interpret an AST with restricted module access."""
     import ast
+
     # Keep only yield-based nodes in fallback.
     unsupported = (ast.Yield, ast.YieldFrom)
     for node in ast.walk(ast_tree):
@@ -756,34 +823,37 @@ def interpret_ast(ast_tree: Any, allowed_modules: list[str], source: str = "") -
                     "float": float,
                     "int": int,
                     "bool": bool,
-                    "Exception": Exception
+                    "Exception": Exception,
                 }
             }
             for mod in allowed_modules:
                 safe_globals[mod] = __import__(mod)
             local_vars = {}
-            # ...existing code...
             exec(compile(ast_tree, "<string>", "exec"), safe_globals, local_vars)
             return local_vars.get("result", None)
     # Otherwise, use the custom interpreter.
     interpreter = ASTInterpreter(allowed_modules=allowed_modules, source=source)
     return interpreter.visit(ast_tree)
 
+
 # A helper function which takes a Python code string and a list of allowed module names,
 # then parses and interprets the code.
 def interpret_code(source_code: str, allowed_modules: List[str]) -> Any:
     """
     Interpret a Python source code string with a restricted set of allowed modules.
-    
-    :param source_code: The Python source code to interpret.
-    :param allowed_modules: A list of module names that are allowed.
-    :return: The result of interpreting the source code.
+
+    Args:
+        source_code: The Python source code to interpret.
+        allowed_modules: A list of module names that are allowed.
+    Returns:
+        The result of interpreting the source code.
     """
     # Dedent the source to normalize its indentation.
     dedented_source = textwrap.dedent(source_code)
     tree: ast.AST = ast.parse(dedented_source)
     return interpret_ast(tree, allowed_modules, source=dedented_source)
 
+
 if __name__ == "__main__":
     print("Script is running!")
     source_code_1: str = """