kailash 0.3.1__py3-none-any.whl → 0.4.0__py3-none-any.whl

kailash/nodes/code/python.py

@@ -112,14 +112,21 @@ class SafeCodeChecker(ast.NodeVisitor):
 
     def __init__(self):
         self.violations = []
+        self.imports_found = []
 
     def visit_Import(self, node):
         """Check import statements."""
         for alias in node.names:
             module_name = alias.name.split(".")[0]
+            self.imports_found.append(module_name)
             if module_name not in ALLOWED_MODULES:
                 self.violations.append(
-                    f"Import of module '{module_name}' is not allowed"
+                    {
+                        "type": "import",
+                        "module": module_name,
+                        "line": node.lineno,
+                        "message": f"Import of module '{module_name}' is not allowed",
+                    }
                 )
         self.generic_visit(node)
 
@@ -127,9 +134,15 @@ class SafeCodeChecker(ast.NodeVisitor):
         """Check from imports."""
         if node.module:
             module_name = node.module.split(".")[0]
+            self.imports_found.append(module_name)
             if module_name not in ALLOWED_MODULES:
                 self.violations.append(
-                    f"Import from module '{module_name}' is not allowed"
+                    {
+                        "type": "import_from",
+                        "module": module_name,
+                        "line": node.lineno,
+                        "message": f"Import from module '{module_name}' is not allowed",
+                    }
                 )
         self.generic_visit(node)
 
@@ -139,12 +152,24 @@ class SafeCodeChecker(ast.NodeVisitor):
             func_name = node.func.id
             # Check for dangerous built-in functions
             if func_name in {"eval", "exec", "compile"}:
-                self.violations.append(f"Call to '{func_name}' is not allowed")
+                self.violations.append(
+                    {
+                        "type": "function_call",
+                        "function": func_name,
+                        "line": node.lineno,
+                        "message": f"Call to '{func_name}' is not allowed",
+                    }
+                )
         elif isinstance(node.func, ast.Attribute):
             # Check for dangerous method calls
             if node.func.attr in {"system", "popen"}:
                 self.violations.append(
-                    f"Call to method '{node.func.attr}' is not allowed"
+                    {
+                        "type": "method_call",
+                        "method": node.func.attr,
+                        "line": node.lineno,
+                        "message": f"Call to method '{node.func.attr}' is not allowed",
+                    }
                 )
         self.generic_visit(node)
 
@@ -249,12 +274,15 @@ class CodeExecutor:
         }
         self._execution_namespace = {}
 
-    def check_code_safety(self, code: str) -> None:
+    def check_code_safety(self, code: str) -> tuple[bool, list[dict], list[str]]:
         """Check if code is safe to execute.
 
         Args:
             code: Python code to check
 
+        Returns:
+            Tuple of (is_safe, violations, imports_found)
+
         Raises:
             SafetyViolationError: If code contains unsafe operations
         """
@@ -264,11 +292,67 @@ class CodeExecutor:
             checker.visit(tree)
 
             if checker.violations:
-                raise SafetyViolationError(
-                    f"Code contains unsafe operations: {'; '.join(checker.violations)}"
-                )
+                # Create detailed error message with suggestions
+                error_parts = []
+                suggestions = []
+
+                for violation in checker.violations:
+                    error_parts.append(
+                        f"Line {violation['line']}: {violation['message']}"
+                    )
+
+                    # Add suggestions based on violation type
+                    if violation["type"] in ["import", "import_from"]:
+                        module = violation["module"]
+                        suggestions.append(
+                            f"Module '{module}' is not allowed. Available modules: {', '.join(sorted(ALLOWED_MODULES))}"
+                        )
+
+                        # Suggest alternatives for common cases
+                        if module == "subprocess":
+                            suggestions.append(
+                                "For file operations, use 'os' or 'pathlib' modules instead"
+                            )
+                        elif module == "requests":
+                            suggestions.append(
+                                "For HTTP requests, use HTTPRequestNode instead of importing requests"
+                            )
+                        elif module == "sqlite3" or module == "psycopg2":
+                            suggestions.append(
+                                "For database operations, use SQLDatabaseNode instead"
+                            )
+                        elif module == "boto3":
+                            suggestions.append(
+                                "For AWS operations, create a custom node or use existing cloud nodes"
+                            )
+
+                    elif violation["type"] == "function_call":
+                        func = violation["function"]
+                        if func in ["eval", "exec"]:
+                            suggestions.append(
+                                f"'{func}' is dangerous. Write explicit code instead of dynamic execution"
+                            )
+                        elif func == "compile":
+                            suggestions.append(
+                                "'compile' is not allowed. Use standard Python code instead"
+                            )
+
+                error_msg = "Code safety violations found:\n" + "\n".join(error_parts)
+                if suggestions:
+                    error_msg += "\n\nSuggestions:\n" + "\n".join(
+                        f"- {s}" for s in suggestions
+                    )
+
+                raise SafetyViolationError(error_msg)
+
+            return True, checker.violations, checker.imports_found
+
         except SyntaxError as e:
-            raise NodeExecutionError(f"Invalid Python syntax: {e}")
+            raise NodeExecutionError(
+                f"Invalid Python syntax at line {e.lineno}: {e.msg}\n"
+                f"Text: {e.text}\n"
+                f"Error position: {' ' * (e.offset - 1) if e.offset else ''}^"
+            )
 
     def execute_code(self, code: str, inputs: dict[str, Any]) -> dict[str, Any]:
         """Execute Python code with given inputs.
@@ -286,7 +370,7 @@ class CodeExecutor:
             MemoryLimitError: If memory usage exceeds limit
         """
         # Check code safety first
-        self.check_code_safety(code)
+        is_safe, violations, imports_found = self.check_code_safety(code)
 
         # Sanitize inputs
         sanitized_inputs = validate_node_parameters(inputs, self.security_config)
@@ -451,9 +535,13 @@ class FunctionWrapper:
         """Execute the wrapped function."""
         result = self.executor.execute_function(self.func, inputs)
 
-        # Wrap non-dict results in a dict
+        # Always wrap results in "result" key for consistent validation
+        # This ensures both dict and non-dict returns have the same structure
         if not isinstance(result, dict):
             result = {"result": result}
+        else:
+            # For dict results, wrap the entire dict in "result" key
+            result = {"result": result}
 
         return result
 
@@ -604,9 +692,13 @@ class ClassWrapper:
         # Execute the method
         result = self.executor.execute_function(method, inputs)
 
-        # Wrap non-dict results in a dict
+        # Always wrap results in "result" key for consistent validation
+        # This ensures both dict and non-dict returns have the same structure
         if not isinstance(result, dict):
             result = {"result": result}
+        else:
+            # For dict results, wrap the entire dict in "result" key
+            result = {"result": result}
 
         return result
 
@@ -717,6 +809,7 @@ class PythonCodeNode(Node):
         input_schema: dict[str, "NodeParameter"] | None = None,
         output_schema: dict[str, "NodeParameter"] | None = None,
         description: str | None = None,
+        max_code_lines: int = 10,
         **kwargs,
     ):
         """Initialize a Python code node.
@@ -732,6 +825,7 @@ class PythonCodeNode(Node):
             input_schema: Explicit input parameter schema for validation
             output_schema: Explicit output parameter schema for validation
             description: Node description
+            max_code_lines: Maximum lines before warning (default: 10)
             **kwargs: Additional node parameters
         """
         # Validate inputs
@@ -753,6 +847,20 @@ class PythonCodeNode(Node):
         self.output_type = output_type or Any
         self._input_schema = input_schema
         self._output_schema = output_schema
+        self.max_code_lines = max_code_lines
+
+        # Check code length and warn if exceeds threshold
+        if self.code and self.max_code_lines > 0:
+            code_lines = [
+                line for line in self.code.strip().split("\n") if line.strip()
+            ]
+            if len(code_lines) > self.max_code_lines:
+                logger.warning(
+                    f"PythonCodeNode '{name}' contains {len(code_lines)} lines of code, "
+                    f"exceeding the recommended maximum of {self.max_code_lines} lines. "
+                    "Consider using PythonCodeNode.from_function() or from_file() for better "
+                    "code organization and IDE support."
+                )
 
         # For class-based nodes, maintain instance
         self.instance = None
@@ -916,6 +1024,32 @@ class PythonCodeNode(Node):
 
         except NodeExecutionError:
             raise
+        except ImportError as e:
+            # Enhanced import error handling
+            module_name = str(e).split("'")[1] if "'" in str(e) else "unknown"
+            error_msg = f"Import error: {str(e)}\n\n"
+
+            # Check if module is in allowed list
+            if module_name not in ALLOWED_MODULES:
+                error_msg += f"Module '{module_name}' is not in the allowed list.\n"
+                error_msg += (
+                    f"Allowed modules: {', '.join(sorted(ALLOWED_MODULES))}\n\n"
+                )
+
+                # Suggest alternatives
+                if module_name == "requests":
+                    error_msg += "Suggestion: Use HTTPRequestNode for HTTP requests instead of importing requests.\n"
+                elif module_name in ["sqlite3", "psycopg2", "pymongo"]:
+                    error_msg += (
+                        "Suggestion: Use SQLDatabaseNode for database operations.\n"
+                    )
+                elif module_name == "boto3":
+                    error_msg += "Suggestion: Create a custom node for AWS operations or use cloud-specific nodes.\n"
+            else:
+                error_msg += f"Module '{module_name}' is allowed but not installed.\n"
+                error_msg += "Suggestion: Install the module using pip or check your environment.\n"
+
+            raise NodeExecutionError(error_msg)
         except Exception as e:
             logger.error(f"Python code execution failed: {e}")
             raise NodeExecutionError(f"Execution failed: {str(e)}")
@@ -1154,3 +1288,150 @@ class PythonCodeNode(Node):
             config["process_method"] = self.process_method
 
         return config
+
+    @staticmethod
+    def list_allowed_modules() -> list[str]:
+        """List all allowed modules for import in PythonCodeNode.
+
+        Returns:
+            Sorted list of allowed module names
+        """
+        return sorted(ALLOWED_MODULES)
+
+    @staticmethod
+    def check_module_availability(module_name: str) -> dict[str, Any]:
+        """Check if a module is allowed and available for import.
+
+        Args:
+            module_name: Name of the module to check
+
+        Returns:
+            Dictionary with status information
+        """
+        result = {
+            "module": module_name,
+            "allowed": module_name in ALLOWED_MODULES,
+            "installed": False,
+            "importable": False,
+            "error": None,
+            "suggestions": [],
+        }
+
+        if not result["allowed"]:
+            result["suggestions"].append(
+                f"Module '{module_name}' is not in the allowed list."
+            )
+            result["suggestions"].append(
+                f"Allowed modules: {', '.join(sorted(ALLOWED_MODULES))}"
+            )
+
+            # Add specific suggestions for common modules
+            if module_name == "requests":
+                result["suggestions"].append(
+                    "Use HTTPRequestNode for HTTP requests instead."
+                )
+            elif module_name in ["sqlite3", "psycopg2", "pymongo", "mysql"]:
+                result["suggestions"].append(
+                    "Use SQLDatabaseNode for database operations."
+                )
+            elif module_name == "boto3":
+                result["suggestions"].append(
+                    "Use cloud-specific nodes or create a custom node."
+                )
+            elif module_name == "subprocess":
+                result["suggestions"].append(
+                    "For security reasons, subprocess is not allowed. Use os or pathlib for file operations."
+                )
+        else:
+            # Check if module is installed
+            try:
+                spec = importlib.util.find_spec(module_name)
+                result["installed"] = spec is not None
+
+                if result["installed"]:
+                    # Try to import it
+                    try:
+                        importlib.import_module(module_name)
+                        result["importable"] = True
+                    except Exception as e:
+                        result["error"] = str(e)
+                        result["suggestions"].append(
+                            f"Module is installed but cannot be imported: {e}"
+                        )
+                else:
+                    result["suggestions"].append(
+                        f"Module '{module_name}' needs to be installed: pip install {module_name}"
+                    )
+            except Exception as e:
+                result["error"] = str(e)
+                result["suggestions"].append(f"Error checking module: {e}")
+
+        return result
+
+    def validate_code(self, code: str) -> dict[str, Any]:
+        """Validate Python code and provide detailed feedback.
+
+        Args:
+            code: Python code to validate
+
+        Returns:
+            Dictionary with validation results
+        """
+        result = {
+            "valid": True,
+            "syntax_errors": [],
+            "safety_violations": [],
+            "imports": [],
+            "suggestions": [],
+            "warnings": [],
+        }
+
+        # Check syntax
+        try:
+            ast.parse(code)
+        except SyntaxError as e:
+            result["valid"] = False
+            result["syntax_errors"].append(
+                {"line": e.lineno, "column": e.offset, "message": e.msg, "text": e.text}
+            )
+            result["suggestions"].append(
+                f"Fix syntax error at line {e.lineno}: {e.msg}"
+            )
+            return result
+
+        # Check safety
+        try:
+            is_safe, violations, imports_found = self.executor.check_code_safety(code)
+            result["imports"] = imports_found
+
+            if violations:
+                result["valid"] = False
+                result["safety_violations"] = violations
+
+                # Add suggestions from violations
+                for violation in violations:
+                    if violation["type"] in ["import", "import_from"]:
+                        module_info = self.check_module_availability(
+                            violation["module"]
+                        )
+                        result["suggestions"].extend(module_info["suggestions"])
+
+        except Exception as e:
+            result["warnings"].append(f"Could not complete safety check: {e}")
+
+        # Check for common issues
+        if "print(" in code and "result" not in code:
+            result["warnings"].append(
+                "Code uses print() but doesn't set 'result'. Output might not be captured."
+            )
+            result["suggestions"].append(
+                "Set 'result' variable to return values from the node."
+            )
+
+        if "input(" in code:
+            result["warnings"].append("Code uses input() which will block execution.")
+            result["suggestions"].append(
+                "Use node parameters instead of input() for user input."
+            )
+
+        return result

kailash/nodes/compliance/__init__.py

@@ -0,0 +1,9 @@
+"""Compliance-related nodes for the Kailash SDK."""
+
+from .data_retention import DataRetentionPolicyNode
+from .gdpr import GDPRComplianceNode
+
+__all__ = [
+    "GDPRComplianceNode",
+    "DataRetentionPolicyNode",
+]