kailash 0.3.0__py3-none-any.whl → 0.3.2__py3-none-any.whl

kailash/runtime/access_controlled.py

@@ -27,7 +27,7 @@ Example with access control (opt-in):
 """
 
 import logging
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any
 
 from kailash.access_control import (
     AccessControlManager,
@@ -108,7 +108,7 @@ class AccessControlledRuntime:
     """
 
     def __init__(
-        self, user_context: UserContext, base_runtime: Optional[LocalRuntime] = None
+        self, user_context: UserContext, base_runtime: LocalRuntime | None = None
     ):
         """
         Initialize access-controlled runtime.
@@ -123,11 +123,11 @@ class AccessControlledRuntime:
 
         # Track skipped nodes for alternative routing
         self._skipped_nodes: set[str] = set()
-        self._node_outputs: Dict[str, Any] = {}
+        self._node_outputs: dict[str, Any] = {}
 
     def execute(
-        self, workflow: Workflow, parameters: Optional[Dict[str, Any]] = None
-    ) -> Tuple[Any, str]:
+        self, workflow: Workflow, parameters: dict[str, Any] | None = None
+    ) -> tuple[Any, str]:
         """
         Execute workflow with access control.
 
@@ -272,7 +272,7 @@ class AccessControlledRuntime:
         return wrapper
 
     @staticmethod
-    def _mask_fields(data: Dict[str, Any], fields: List[str]) -> Dict[str, Any]:
+    def _mask_fields(data: dict[str, Any], fields: list[str]) -> dict[str, Any]:
         """Mask sensitive fields in data"""
         masked = data.copy()
         for field in fields:
@@ -281,8 +281,8 @@ class AccessControlledRuntime:
         return masked
 
     def _handle_conditional_routing(
-        self, node_id: str, true_path: List[str], false_path: List[str]
-    ) -> List[str]:
+        self, node_id: str, true_path: list[str], false_path: list[str]
+    ) -> list[str]:
         """
         Determine which path to take based on permissions.
 
@@ -345,14 +345,14 @@ class AccessControlConfig:
     """
 
     def __init__(self):
-        self.rules: List[PermissionRule] = []
+        self.rules: list[PermissionRule] = []
 
     def add_workflow_permission(
         self,
         workflow_id: str,
         permission: WorkflowPermission,
-        user_id: Optional[str] = None,
-        role: Optional[str] = None,
+        user_id: str | None = None,
+        role: str | None = None,
         effect: PermissionEffect = PermissionEffect.ALLOW,
     ):
         """Add a workflow-level permission rule"""
@@ -372,11 +372,11 @@ class AccessControlConfig:
         workflow_id: str,
         node_id: str,
         permission: NodePermission,
-        user_id: Optional[str] = None,
-        role: Optional[str] = None,
+        user_id: str | None = None,
+        role: str | None = None,
         effect: PermissionEffect = PermissionEffect.ALLOW,
-        masked_fields: Optional[List[str]] = None,
-        redirect_node: Optional[str] = None,
+        masked_fields: list[str] | None = None,
+        redirect_node: str | None = None,
     ):
         """Add a node-level permission rule"""
         rule = PermissionRule(
@@ -406,9 +406,9 @@ class AccessControlConfig:
 def execute_with_access_control(
     workflow: Workflow,
     user_context: UserContext,
-    parameters: Optional[Dict[str, Any]] = None,
-    access_config: Optional[AccessControlConfig] = None,
-) -> Tuple[Any, str]:
+    parameters: dict[str, Any] | None = None,
+    access_config: AccessControlConfig | None = None,
+) -> tuple[Any, str]:
     """
     Convenience function to execute a workflow with access control.
 

kailash/runtime/async_local.py

@@ -6,8 +6,8 @@ database queries, or LLM interactions.
 """
 
 import logging
-from datetime import datetime, timezone
-from typing import Any, Dict, Optional, Tuple
+from datetime import UTC, datetime
+from typing import Any
 
 import networkx as nx
 
@@ -60,9 +60,9 @@ class AsyncLocalRuntime:
     async def execute(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager] = None,
-        parameters: Optional[Dict[str, Dict[str, Any]]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None = None,
+        parameters: dict[str, dict[str, Any]] | None = None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a workflow asynchronously.
 
         Args:
@@ -144,10 +144,10 @@ class AsyncLocalRuntime:
     async def _execute_workflow(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-        parameters: Dict[str, Dict[str, Any]],
-    ) -> Dict[str, Any]:
+        task_manager: TaskManager | None,
+        run_id: str | None,
+        parameters: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
         """Execute the workflow nodes asynchronously.
 
         Args:
@@ -195,7 +195,7 @@ class AsyncLocalRuntime:
                         run_id=run_id,
                         node_id=node_id,
                         node_type=node_instance.__class__.__name__,
-                        started_at=datetime.now(timezone.utc),
+                        started_at=datetime.now(UTC),
                     )
                 except Exception as e:
                     self.logger.warning(
@@ -220,7 +220,7 @@ class AsyncLocalRuntime:
                     task.update_status(TaskStatus.RUNNING)
 
                 # Execute node - check if it supports async execution
-                start_time = datetime.now(timezone.utc)
+                start_time = datetime.now(UTC)
 
                 if isinstance(node_instance, AsyncNode):
                     # Use async execution
@@ -229,9 +229,7 @@ class AsyncLocalRuntime:
                     # Fall back to synchronous execution
                     outputs = node_instance.run(**inputs)
 
-                execution_time = (
-                    datetime.now(timezone.utc) - start_time
-                ).total_seconds()
+                execution_time = (datetime.now(UTC) - start_time).total_seconds()
 
                 # Store outputs
                 node_outputs[node_id] = outputs
@@ -245,7 +243,7 @@ class AsyncLocalRuntime:
                     task.update_status(
                         TaskStatus.COMPLETED,
                         result=outputs,
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                         metadata={"execution_time": execution_time},
                     )
 
@@ -262,7 +260,7 @@ class AsyncLocalRuntime:
                     task.update_status(
                         TaskStatus.FAILED,
                         error=str(e),
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                     )
 
                 # Determine if we should continue or stop
@@ -287,9 +285,9 @@ class AsyncLocalRuntime:
         workflow: Workflow,
         node_id: str,
         node_instance: Any,
-        node_outputs: Dict[str, Dict[str, Any]],
-        parameters: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        node_outputs: dict[str, dict[str, Any]],
+        parameters: dict[str, Any],
+    ) -> dict[str, Any]:
         """Prepare inputs for a node execution.
 
         Args:

kailash/runtime/docker.py

@@ -21,7 +21,7 @@ import subprocess
 import sys
 import tempfile
 from pathlib import Path
-from typing import Any, Dict, Optional, Tuple
+from typing import Any
 
 from kailash.nodes.base import Node
 
@@ -50,8 +50,8 @@ class DockerNodeWrapper:
         node: Node,
         node_id: str,
         base_image: str = "python:3.11-slim",
-        work_dir: Optional[Path] = None,
-        sdk_path: Optional[Path] = None,
+        work_dir: Path | None = None,
+        sdk_path: Path | None = None,
     ):
         """
         Initialize a Docker node wrapper.
@@ -335,7 +335,7 @@ ENTRYPOINT ["/app/entrypoint.py"]
             logger.error(error_msg)
             raise RuntimeError(error_msg)
 
-    def prepare_inputs(self, inputs: Dict[str, Any]):
+    def prepare_inputs(self, inputs: dict[str, Any]):
         """
         Prepare inputs for node execution.
 
@@ -349,8 +349,8 @@ ENTRYPOINT ["/app/entrypoint.py"]
     def run_container(
         self,
         network: str = None,
-        env_vars: Dict[str, str] = None,
-        resource_limits: Dict[str, str] = None,
+        env_vars: dict[str, str] = None,
+        resource_limits: dict[str, str] = None,
     ) -> bool:
         """
         Run the node in a Docker container.
@@ -416,13 +416,13 @@ ENTRYPOINT ["/app/entrypoint.py"]
             # Check if there's an error file
             error_file = self.output_dir / "error.json"
             if error_file.exists():
-                with open(error_file, "r") as f:
+                with open(error_file) as f:
                     error_data = json.load(f)
                     error_msg = f"Node execution error: {error_data.get('error', 'Unknown error')}"
 
             raise NodeExecutionError(error_msg)
 
-    def get_results(self) -> Dict[str, Any]:
+    def get_results(self) -> dict[str, Any]:
         """
         Get the results of node execution.
 
@@ -431,12 +431,12 @@ ENTRYPOINT ["/app/entrypoint.py"]
         """
         result_file = self.output_dir / "result.json"
         if result_file.exists():
-            with open(result_file, "r") as f:
+            with open(result_file) as f:
                 return json.load(f)
 
         error_file = self.output_dir / "error.json"
         if error_file.exists():
-            with open(error_file, "r") as f:
+            with open(error_file) as f:
                 error_data = json.load(f)
                 raise NodeExecutionError(
                     f"Node {self.node_id} execution failed: {error_data.get('error', 'Unknown error')}"
@@ -464,10 +464,10 @@ class DockerRuntime:
         self,
         base_image: str = "python:3.11-slim",
         network_name: str = "kailash-network",
-        work_dir: Optional[str] = None,
-        sdk_path: Optional[str] = None,
-        resource_limits: Optional[Dict[str, str]] = None,
-        task_manager: Optional[TaskManager] = None,
+        work_dir: str | None = None,
+        sdk_path: str | None = None,
+        resource_limits: dict[str, str] | None = None,
+        task_manager: TaskManager | None = None,
     ):
         """
         Initialize the Docker runtime.
@@ -516,14 +516,14 @@ class DockerRuntime:
         # Track node wrappers
         self.node_wrappers = {}
 
-    def _create_task_run(self, workflow: Workflow) -> Optional[str]:
+    def _create_task_run(self, workflow: Workflow) -> str | None:
         """Create a task run if task manager is available."""
         if self.task_manager:
             return self.task_manager.create_run(workflow.name)
         return None
 
     def _update_task_status(
-        self, run_id: Optional[str], node_id: str, status: str, output: Any = None
+        self, run_id: str | None, node_id: str, status: str, output: Any = None
     ):
         """Update task status if task manager is available."""
         if self.task_manager and run_id:
@@ -534,9 +534,7 @@ class DockerRuntime:
                 )
                 self.task_manager.update_run_status(run_id, "failed", error_msg)
 
-    def _complete_task_run(
-        self, run_id: Optional[str], status: str, result: Any = None
-    ):
+    def _complete_task_run(self, run_id: str | None, status: str, result: Any = None):
         """Complete task run if task manager is available."""
         if self.task_manager and run_id:
             if status == "completed":
@@ -567,9 +565,9 @@ class DockerRuntime:
     def execute(
         self,
         workflow: Workflow,
-        inputs: Dict[str, Dict[str, Any]] = None,
-        node_resource_limits: Dict[str, Dict[str, str]] = None,
-    ) -> Tuple[Dict[str, Dict[str, Any]], str]:
+        inputs: dict[str, dict[str, Any]] = None,
+        node_resource_limits: dict[str, dict[str, str]] = None,
+    ) -> tuple[dict[str, dict[str, Any]], str]:
         """
         Execute a workflow using Docker containers.
 

kailash/runtime/local.py

@@ -40,8 +40,8 @@ Examples:
 """
 
 import logging
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Tuple
+from datetime import UTC, datetime
+from typing import Any
 
 import networkx as nx
 
@@ -90,9 +90,9 @@ class LocalRuntime:
     def execute(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager] = None,
-        parameters: Optional[Dict[str, Dict[str, Any]]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None = None,
+        parameters: dict[str, dict[str, Any]] | None = None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a workflow locally.
 
         Args:
@@ -197,10 +197,10 @@ class LocalRuntime:
     def _execute_workflow(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-        parameters: Dict[str, Dict[str, Any]],
-    ) -> Dict[str, Any]:
+        task_manager: TaskManager | None,
+        run_id: str | None,
+        parameters: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
         """Execute the workflow nodes in topological order.
 
         Args:
@@ -273,7 +273,7 @@ class LocalRuntime:
                         run_id=run_id,
                         node_id=node_id,
                         node_type=node_instance.__class__.__name__,
-                        started_at=datetime.now(timezone.utc),
+                        started_at=datetime.now(UTC),
                         metadata=node_metadata,
                     )
                     # Start the task
@@ -329,7 +329,7 @@ class LocalRuntime:
                         task.task_id,
                         TaskStatus.COMPLETED,
                         result=outputs,
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                         metadata={"execution_time": performance_metrics.duration},
                     )
 
@@ -350,7 +350,7 @@ class LocalRuntime:
                         task.task_id,
                         TaskStatus.FAILED,
                         error=str(e),
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                     )
 
                 # Determine if we should continue
@@ -375,9 +375,9 @@ class LocalRuntime:
         workflow: Workflow,
         node_id: str,
         node_instance: Node,
-        node_outputs: Dict[str, Dict[str, Any]],
-        parameters: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        node_outputs: dict[str, dict[str, Any]],
+        parameters: dict[str, Any],
+    ) -> dict[str, Any]:
         """Prepare inputs for a node execution.
 
         Args:
@@ -476,7 +476,7 @@ class LocalRuntime:
         # Future: implement configurable error handling policies
         return has_dependents
 
-    def validate_workflow(self, workflow: Workflow) -> List[str]:
+    def validate_workflow(self, workflow: Workflow) -> list[str]:
         """Validate a workflow before execution.
 
         Args:

kailash/runtime/parallel.py

@@ -8,8 +8,8 @@ import asyncio
 import logging
 import time
 from collections import deque
-from datetime import datetime, timezone
-from typing import Any, Deque, Dict, Optional, Set, Tuple
+from datetime import UTC, datetime
+from typing import Any
 
 import networkx as nx
 
@@ -66,9 +66,9 @@ class ParallelRuntime:
     async def execute(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager] = None,
-        parameters: Optional[Dict[str, Dict[str, Any]]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None = None,
+        parameters: dict[str, dict[str, Any]] | None = None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a workflow with parallel node execution.
 
         Args:
@@ -159,10 +159,10 @@ class ParallelRuntime:
     async def _execute_workflow_parallel(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-        parameters: Dict[str, Dict[str, Any]],
-    ) -> Dict[str, Any]:
+        task_manager: TaskManager | None,
+        run_id: str | None,
+        parameters: dict[str, dict[str, Any]],
+    ) -> dict[str, Any]:
         """Execute the workflow nodes in parallel where possible.
 
         This method uses a dynamic scheduling approach to run independent nodes
@@ -323,11 +323,11 @@ class ParallelRuntime:
         self,
         workflow: Workflow,
         node_id: str,
-        node_outputs: Dict[str, Dict[str, Any]],
-        parameters: Dict[str, Any],
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-    ) -> Tuple[Dict[str, Any], bool]:
+        node_outputs: dict[str, dict[str, Any]],
+        parameters: dict[str, Any],
+        task_manager: TaskManager | None,
+        run_id: str | None,
+    ) -> tuple[dict[str, Any], bool]:
         """Execute a single node asynchronously.
 
         Args:
@@ -359,7 +359,7 @@ class ParallelRuntime:
                     run_id=run_id,
                     node_id=node_id,
                     node_type=node_instance.__class__.__name__,
-                    started_at=datetime.now(timezone.utc),
+                    started_at=datetime.now(UTC),
                 )
         except Exception as e:
             self.logger.warning(f"Failed to create task for node '{node_id}': {e}")
@@ -409,7 +409,7 @@ class ParallelRuntime:
                     task.update_status(
                         TaskStatus.COMPLETED,
                         result=outputs,
-                        ended_at=datetime.now(timezone.utc),
+                        ended_at=datetime.now(UTC),
                         metadata={"execution_time": performance_metrics.duration},
                     )
 
@@ -431,7 +431,7 @@ class ParallelRuntime:
             # Update task status
             if task:
                 task.update_status(
-                    TaskStatus.FAILED, error=str(e), ended_at=datetime.now(timezone.utc)
+                    TaskStatus.FAILED, error=str(e), ended_at=datetime.now(UTC)
                 )
 
             # Return error result
@@ -448,9 +448,9 @@ class ParallelRuntime:
         workflow: Workflow,
         node_id: str,
         node_instance: Any,
-        node_outputs: Dict[str, Dict[str, Any]],
-        parameters: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        node_outputs: dict[str, dict[str, Any]],
+        parameters: dict[str, Any],
+    ) -> dict[str, Any]:
         """Prepare inputs for a node execution.
 
         Args:
@@ -520,9 +520,9 @@ class ParallelRuntime:
         self,
         workflow: Workflow,
         failed_node: str,
-        failed_nodes: Set[str],
-        pending_nodes: Set[str],
-        ready_nodes: Deque[str],
+        failed_nodes: set[str],
+        pending_nodes: set[str],
+        ready_nodes: deque[str],
     ) -> None:
         """Mark all dependent nodes as failed.
 

kailash/runtime/parallel_cyclic.py

@@ -2,8 +2,8 @@
 
 import logging
 from concurrent.futures import ThreadPoolExecutor, as_completed
-from datetime import datetime, timezone
-from typing import Any, Dict, List, Optional, Set, Tuple
+from datetime import UTC, datetime
+from typing import Any
 
 import networkx as nx
 
@@ -56,10 +56,10 @@ class ParallelCyclicRuntime:
     def execute(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager] = None,
-        parameters: Optional[Dict[str, Dict[str, Any]]] = None,
-        parallel_nodes: Optional[Set[str]] = None,
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None = None,
+        parameters: dict[str, dict[str, Any]] | None = None,
+        parallel_nodes: set[str] | None = None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a workflow with parallel and cyclic support.
 
         Args:
@@ -108,9 +108,9 @@ class ParallelCyclicRuntime:
     def _execute_cyclic_workflow(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        parameters: Optional[Dict[str, Dict[str, Any]]],
-    ) -> Tuple[Dict[str, Any], str]:
+        task_manager: TaskManager | None,
+        parameters: dict[str, dict[str, Any]] | None,
+    ) -> tuple[dict[str, Any], str]:
         """Execute a cyclic workflow with potential parallel optimizations.
 
         Args:
@@ -141,10 +141,10 @@ class ParallelCyclicRuntime:
     def _execute_parallel_dag(
         self,
         workflow: Workflow,
-        task_manager: Optional[TaskManager],
-        parameters: Optional[Dict[str, Dict[str, Any]]],
-        parallel_nodes: Optional[Set[str]],
-    ) -> Tuple[Dict[str, Any], Optional[str]]:
+        task_manager: TaskManager | None,
+        parameters: dict[str, dict[str, Any]] | None,
+        parallel_nodes: set[str] | None,
+    ) -> tuple[dict[str, Any], str | None]:
         """Execute a DAG workflow with parallel node execution.
 
         Args:
@@ -252,8 +252,8 @@ class ParallelCyclicRuntime:
             raise
 
     def _analyze_parallel_groups(
-        self, workflow: Workflow, parallel_nodes: Optional[Set[str]]
-    ) -> List[List[str]]:
+        self, workflow: Workflow, parallel_nodes: set[str] | None
+    ) -> list[list[str]]:
         """Analyze workflow to identify groups of nodes that can be executed in parallel.
 
         Args:
@@ -320,11 +320,11 @@ class ParallelCyclicRuntime:
         self,
         workflow: Workflow,
         node_id: str,
-        previous_results: Dict[str, Any],
-        parameters: Optional[Dict[str, Dict[str, Any]]],
-        task_manager: Optional[TaskManager],
-        run_id: Optional[str],
-    ) -> Dict[str, Any]:
+        previous_results: dict[str, Any],
+        parameters: dict[str, dict[str, Any]] | None,
+        task_manager: TaskManager | None,
+        run_id: str | None,
+    ) -> dict[str, Any]:
         """Execute a single node in isolation.
 
         Args:
@@ -356,7 +356,7 @@ class ParallelCyclicRuntime:
                     run_id=run_id,
                     node_id=node_id,
                     node_type=node_instance.__class__.__name__,
-                    started_at=datetime.now(timezone.utc),
+                    started_at=datetime.now(UTC),
                     metadata={},
                 )
                 if task:
@@ -397,7 +397,7 @@ class ParallelCyclicRuntime:
                     task.task_id,
                     TaskStatus.COMPLETED,
                     result=outputs,
-                    ended_at=datetime.now(timezone.utc),
+                    ended_at=datetime.now(UTC),
                     metadata={"execution_time": performance_metrics.duration},
                 )
                 task_manager.update_task_metrics(task.task_id, task_metrics)
@@ -415,7 +415,7 @@ class ParallelCyclicRuntime:
                     task.task_id,
                     TaskStatus.FAILED,
                     error=str(e),
-                    ended_at=datetime.now(timezone.utc),
+                    ended_at=datetime.now(UTC),
                 )
 
             self.logger.error(f"Node {node_id} failed: {e}", exc_info=self.debug)
@@ -428,9 +428,9 @@ class ParallelCyclicRuntime:
         workflow: Workflow,
         node_id: str,
         node_instance: Node,
-        previous_results: Dict[str, Any],
-        parameters: Dict[str, Any],
-    ) -> Dict[str, Any]:
+        previous_results: dict[str, Any],
+        parameters: dict[str, Any],
+    ) -> dict[str, Any]:
         """Prepare inputs for a node execution in parallel context.
 
         Args:
@@ -509,7 +509,7 @@ class ParallelCyclicRuntime:
             return False
 
     def _should_stop_on_group_error(
-        self, workflow: Workflow, failed_node: str, node_group: List[str]
+        self, workflow: Workflow, failed_node: str, node_group: list[str]
     ) -> bool:
         """Determine if execution should stop when a node in a parallel group fails.