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

kailash/tracking/storage/database.py

@@ -2,7 +2,6 @@
 
 import json
 from datetime import datetime
-from typing import List, Optional
 from uuid import uuid4
 
 from ..models import TaskMetrics, TaskRun, TaskStatus, WorkflowRun
@@ -160,7 +159,7 @@ class DatabaseStorage(StorageBackend):
 
         self.conn.commit()
 
-    def load_run(self, run_id: str) -> Optional[WorkflowRun]:
+    def load_run(self, run_id: str) -> WorkflowRun | None:
         """Load a workflow run by ID."""
         cursor = self.conn.cursor()
 
@@ -177,7 +176,7 @@ class DatabaseStorage(StorageBackend):
 
         # Convert row to dict
         columns = [desc[0] for desc in cursor.description]
-        data = dict(zip(columns, row))
+        data = dict(zip(columns, row, strict=False))
 
         # Parse JSON metadata
         data["metadata"] = json.loads(data["metadata"] or "{}")
@@ -189,8 +188,8 @@ class DatabaseStorage(StorageBackend):
         return WorkflowRun.model_validate(data)
 
     def list_runs(
-        self, workflow_name: Optional[str] = None, status: Optional[str] = None
-    ) -> List[WorkflowRun]:
+        self, workflow_name: str | None = None, status: str | None = None
+    ) -> list[WorkflowRun]:
         """List workflow runs."""
         cursor = self.conn.cursor()
 
@@ -213,7 +212,7 @@ class DatabaseStorage(StorageBackend):
         columns = [desc[0] for desc in cursor.description]
 
         for row in cursor.fetchall():
-            data = dict(zip(columns, row))
+            data = dict(zip(columns, row, strict=False))
             data["metadata"] = json.loads(data["metadata"] or "{}")
 
             # Load task IDs
@@ -277,7 +276,7 @@ class DatabaseStorage(StorageBackend):
 
         self.conn.commit()
 
-    def load_task(self, task_id: str) -> Optional[TaskRun]:
+    def load_task(self, task_id: str) -> TaskRun | None:
         """Load a task by ID."""
         cursor = self.conn.cursor()
 
@@ -305,7 +304,7 @@ class DatabaseStorage(StorageBackend):
 
         # Convert row to dict
         columns = [desc[0] for desc in cursor.description]
-        data = dict(zip(columns, row))
+        data = dict(zip(columns, row, strict=False))
 
         # Parse JSON fields
         if data["result"]:
@@ -341,7 +340,7 @@ class DatabaseStorage(StorageBackend):
         metrics_row = cursor.fetchone()
         if metrics_row:
             metrics_columns = [desc[0] for desc in cursor.description]
-            metrics_data = dict(zip(metrics_columns, metrics_row))
+            metrics_data = dict(zip(metrics_columns, metrics_row, strict=False))
 
             # Parse custom metrics if present
             if metrics_data.get("custom_metrics"):
@@ -362,9 +361,9 @@ class DatabaseStorage(StorageBackend):
     def list_tasks(
         self,
         run_id: str,
-        node_id: Optional[str] = None,
-        status: Optional[TaskStatus] = None,
-    ) -> List[TaskRun]:
+        node_id: str | None = None,
+        status: TaskStatus | None = None,
+    ) -> list[TaskRun]:
         """List tasks for a run."""
         cursor = self.conn.cursor()
 
@@ -387,7 +386,7 @@ class DatabaseStorage(StorageBackend):
         columns = [desc[0] for desc in cursor.description]
 
         for row in cursor.fetchall():
-            data = dict(zip(columns, row))
+            data = dict(zip(columns, row, strict=False))
 
             # Parse JSON fields
             if data["result"]:
@@ -451,7 +450,7 @@ class DatabaseStorage(StorageBackend):
 
     def import_run(self, input_path: str) -> str:
         """Import a run and its tasks."""
-        with open(input_path, "r") as f:
+        with open(input_path) as f:
             import_data = json.load(f)
 
         # Import run
@@ -472,7 +471,7 @@ class DatabaseStorage(StorageBackend):
 
         return run.run_id
 
-    def get_task(self, task_id: str) -> Optional[TaskRun]:
+    def get_task(self, task_id: str) -> TaskRun | None:
         """Load a task by ID.
 
         Alias for load_task for API compatibility.
@@ -485,7 +484,7 @@ class DatabaseStorage(StorageBackend):
         """
         return self.load_task(task_id)
 
-    def get_all_tasks(self) -> List[TaskRun]:
+    def get_all_tasks(self) -> list[TaskRun]:
         """Get all tasks.
 
         Returns:
@@ -498,7 +497,7 @@ class DatabaseStorage(StorageBackend):
         columns = [desc[0] for desc in cursor.description]
 
         for row in cursor.fetchall():
-            data = dict(zip(columns, row))
+            data = dict(zip(columns, row, strict=False))
 
             # Parse JSON fields
             if data["result"]:
@@ -532,11 +531,11 @@ class DatabaseStorage(StorageBackend):
 
     def query_tasks(
         self,
-        node_id: Optional[str] = None,
-        status: Optional[TaskStatus] = None,
-        started_after: Optional[datetime] = None,
-        completed_before: Optional[datetime] = None,
-    ) -> List[TaskRun]:
+        node_id: str | None = None,
+        status: TaskStatus | None = None,
+        started_after: datetime | None = None,
+        completed_before: datetime | None = None,
+    ) -> list[TaskRun]:
         """Query tasks with filters.
 
         Args:
@@ -583,7 +582,7 @@ class DatabaseStorage(StorageBackend):
         columns = [desc[0] for desc in cursor.description]
 
         for row in cursor.fetchall():
-            data = dict(zip(columns, row))
+            data = dict(zip(columns, row, strict=False))
 
             # Parse JSON fields
             if data["result"]:

kailash/tracking/storage/filesystem.py

@@ -2,9 +2,9 @@
 
 import json
 import os
-from datetime import datetime
+from datetime import UTC, datetime
 from pathlib import Path
-from typing import Any, List, Optional
+from typing import Any
 from uuid import uuid4
 
 from kailash.sdk_exceptions import KailashStorageError
@@ -16,7 +16,7 @@ from .base import StorageBackend
 class FileSystemStorage(StorageBackend):
     """Filesystem-based storage backend."""
 
-    def __init__(self, base_path: Optional[str] = None):
+    def __init__(self, base_path: str | None = None):
         """Initialize filesystem storage.
 
         Args:
@@ -47,26 +47,26 @@ class FileSystemStorage(StorageBackend):
         with open(run_path, "w") as f:
             json.dump(run.to_dict(), f, indent=2)
 
-    def load_run(self, run_id: str) -> Optional[WorkflowRun]:
+    def load_run(self, run_id: str) -> WorkflowRun | None:
         """Load a workflow run by ID."""
         run_path = self.runs_dir / f"{run_id}.json"
         if not run_path.exists():
             return None
 
-        with open(run_path, "r") as f:
+        with open(run_path) as f:
             data = json.load(f)
 
         return WorkflowRun.model_validate(data)
 
     def list_runs(
-        self, workflow_name: Optional[str] = None, status: Optional[str] = None
-    ) -> List[WorkflowRun]:
+        self, workflow_name: str | None = None, status: str | None = None
+    ) -> list[WorkflowRun]:
         """List workflow runs."""
         runs = []
 
         for run_file in self.runs_dir.glob("*.json"):
             try:
-                with open(run_file, "r") as f:
+                with open(run_file) as f:
                     data = json.load(f)
 
                 run = WorkflowRun.model_validate(data)
@@ -89,14 +89,13 @@ class FileSystemStorage(StorageBackend):
                 # Ensure datetime is timezone-aware
                 if run.started_at.tzinfo is None:
                     # Assume UTC for naive datetimes
-                    from datetime import timezone
 
-                    return run.started_at.replace(tzinfo=timezone.utc)
+                    return run.started_at.replace(tzinfo=UTC)
                 return run.started_at
             # Return a very old date for runs without started_at
-            from datetime import datetime, timezone
+            from datetime import datetime
 
-            return datetime.min.replace(tzinfo=timezone.utc)
+            return datetime.min.replace(tzinfo=UTC)
 
         runs.sort(key=safe_datetime_key, reverse=True)
         return runs
@@ -135,7 +134,7 @@ class FileSystemStorage(StorageBackend):
         except Exception as e:
             raise KailashStorageError(f"Failed to save task: {e}") from e
 
-    def get_task(self, task_id: str) -> Optional[TaskRun]:
+    def get_task(self, task_id: str) -> TaskRun | None:
         """Load a task by ID.
 
         Args:
@@ -151,14 +150,14 @@ class FileSystemStorage(StorageBackend):
             # First check direct path for tests
             task_path = self.tasks_dir / f"{task_id}.json"
             if task_path.exists():
-                with open(task_path, "r") as tf:
+                with open(task_path) as tf:
                     task_data = json.load(tf)
                 task = TaskRun.model_validate(task_data)
 
                 # Load metrics if available
                 metrics_path = self.metrics_dir / f"{task_id}.json"
                 if metrics_path.exists():
-                    with open(metrics_path, "r") as mf:
+                    with open(metrics_path) as mf:
                         metrics_data = json.load(mf)
                     task.metrics = TaskMetrics.model_validate(metrics_data)
 
@@ -167,20 +166,20 @@ class FileSystemStorage(StorageBackend):
             # Then check index for run_id
             index_path = self._get_index_file()
             if index_path.exists():
-                with open(index_path, "r") as f:
+                with open(index_path) as f:
                     index = json.load(f)
                     if task_id in index.get("tasks", {}):
                         run_id = index["tasks"][task_id]["run_id"]
                         run_task_path = self.tasks_dir / run_id / f"{task_id}.json"
                         if run_task_path.exists():
-                            with open(run_task_path, "r") as tf:
+                            with open(run_task_path) as tf:
                                 task_data = json.load(tf)
                             task = TaskRun.model_validate(task_data)
 
                             # Load metrics if available
                             metrics_path = self.metrics_dir / f"{task_id}.json"
                             if metrics_path.exists():
-                                with open(metrics_path, "r") as mf:
+                                with open(metrics_path) as mf:
                                     metrics_data = json.load(mf)
                                 task.metrics = TaskMetrics.model_validate(metrics_data)
 
@@ -193,7 +192,7 @@ class FileSystemStorage(StorageBackend):
                 raise
             raise KailashStorageError(f"Failed to get task: {e}") from e
 
-    def load_task(self, task_id: str) -> Optional[TaskRun]:
+    def load_task(self, task_id: str) -> TaskRun | None:
         """Load a task by ID."""
         # Search all run directories
         for run_dir in self.tasks_dir.iterdir():
@@ -202,7 +201,7 @@ class FileSystemStorage(StorageBackend):
 
             task_path = run_dir / f"{task_id}.json"
             if task_path.exists():
-                with open(task_path, "r") as f:
+                with open(task_path) as f:
                     data = json.load(f)
                 return TaskRun.model_validate(data)
 
@@ -211,9 +210,9 @@ class FileSystemStorage(StorageBackend):
     def list_tasks(
         self,
         run_id: str,
-        node_id: Optional[str] = None,
-        status: Optional[TaskStatus] = None,
-    ) -> List[TaskRun]:
+        node_id: str | None = None,
+        status: TaskStatus | None = None,
+    ) -> list[TaskRun]:
         """List tasks for a run."""
         tasks = []
         run_tasks_dir = self.tasks_dir / run_id
@@ -223,7 +222,7 @@ class FileSystemStorage(StorageBackend):
 
         for task_file in run_tasks_dir.glob("*.json"):
             try:
-                with open(task_file, "r") as f:
+                with open(task_file) as f:
                     data = json.load(f)
 
                 task = TaskRun.model_validate(data)
@@ -246,9 +245,8 @@ class FileSystemStorage(StorageBackend):
                 # Ensure datetime is timezone-aware
                 if task.started_at.tzinfo is None:
                     # Assume UTC for naive datetimes
-                    from datetime import timezone
 
-                    return task.started_at.replace(tzinfo=timezone.utc)
+                    return task.started_at.replace(tzinfo=UTC)
                 return task.started_at
             # Use task_id as fallback for tasks without started_at
             return task.task_id
@@ -290,7 +288,7 @@ class FileSystemStorage(StorageBackend):
 
     def import_run(self, input_path: str) -> str:
         """Import a run and its tasks."""
-        with open(input_path, "r") as f:
+        with open(input_path) as f:
             import_data = json.load(f)
 
         # Import run
@@ -383,7 +381,7 @@ class FileSystemStorage(StorageBackend):
             # Update index
             index_path = self._get_index_file()
             if index_path.exists():
-                with open(index_path, "r") as f:
+                with open(index_path) as f:
                     index = json.load(f)
 
                 if task_id in index.get("tasks", {}):
@@ -396,7 +394,7 @@ class FileSystemStorage(StorageBackend):
                 raise
             raise KailashStorageError(f"Failed to delete task: {e}") from e
 
-    def get_all_tasks(self) -> List[TaskRun]:
+    def get_all_tasks(self) -> list[TaskRun]:
         """Get all tasks.
 
         Returns:
@@ -411,7 +409,7 @@ class FileSystemStorage(StorageBackend):
             # First load tasks in the main tasks directory (for tests)
             for task_file in self.tasks_dir.glob("*.json"):
                 if task_file.is_file():
-                    with open(task_file, "r") as f:
+                    with open(task_file) as f:
                         task_data = json.load(f)
 
                     task = TaskRun.model_validate(task_data)
@@ -419,7 +417,7 @@ class FileSystemStorage(StorageBackend):
                     # Load metrics if available
                     metrics_path = self.metrics_dir / f"{task.task_id}.json"
                     if metrics_path.exists():
-                        with open(metrics_path, "r") as f:
+                        with open(metrics_path) as f:
                             metrics_data = json.load(f)
                         task.metrics = TaskMetrics.model_validate(metrics_data)
 
@@ -432,7 +430,7 @@ class FileSystemStorage(StorageBackend):
 
                 # Load all tasks in the run directory
                 for task_file in run_dir.glob("*.json"):
-                    with open(task_file, "r") as f:
+                    with open(task_file) as f:
                         task_data = json.load(f)
 
                     task = TaskRun.model_validate(task_data)
@@ -440,7 +438,7 @@ class FileSystemStorage(StorageBackend):
                     # Load metrics if available
                     metrics_path = self.metrics_dir / f"{task.task_id}.json"
                     if metrics_path.exists():
-                        with open(metrics_path, "r") as f:
+                        with open(metrics_path) as f:
                             metrics_data = json.load(f)
                         task.metrics = TaskMetrics.model_validate(metrics_data)
 
@@ -450,7 +448,7 @@ class FileSystemStorage(StorageBackend):
         except Exception as e:
             raise KailashStorageError(f"Failed to get all tasks: {e}") from e
 
-    def get_tasks_by_run(self, run_id: str) -> List[TaskRun]:
+    def get_tasks_by_run(self, run_id: str) -> list[TaskRun]:
         """Get all tasks for a specific run.
 
         Args:
@@ -466,11 +464,11 @@ class FileSystemStorage(StorageBackend):
 
     def query_tasks(
         self,
-        node_id: Optional[str] = None,
-        status: Optional[TaskStatus] = None,
-        started_after: Optional[datetime] = None,
-        completed_before: Optional[datetime] = None,
-    ) -> List[TaskRun]:
+        node_id: str | None = None,
+        status: TaskStatus | None = None,
+        started_after: datetime | None = None,
+        completed_before: datetime | None = None,
+    ) -> list[TaskRun]:
         """Query tasks with filters.
 
         Args:
@@ -536,7 +534,7 @@ class FileSystemStorage(StorageBackend):
             # Load existing index
             if index_path.exists():
                 try:
-                    with open(index_path, "r") as f:
+                    with open(index_path) as f:
                         index = json.load(f)
                 except json.JSONDecodeError:
                     # Handle case where the file is empty or invalid

kailash/utils/export.py

@@ -5,7 +5,7 @@ import logging
 import re
 from copy import deepcopy
 from pathlib import Path
-from typing import Any, Dict, List, Optional, Set
+from typing import Any
 
 import yaml
 from pydantic import BaseModel, Field, ValidationError
@@ -26,9 +26,9 @@ class ResourceSpec(BaseModel):
 
     cpu: str = Field("100m", description="CPU request")
     memory: str = Field("128Mi", description="Memory request")
-    cpu_limit: Optional[str] = Field(None, description="CPU limit")
-    memory_limit: Optional[str] = Field(None, description="Memory limit")
-    gpu: Optional[int] = Field(None, description="Number of GPUs")
+    cpu_limit: str | None = Field(None, description="CPU limit")
+    memory_limit: str | None = Field(None, description="Memory limit")
+    gpu: int | None = Field(None, description="Number of GPUs")
 
 
 class ContainerMapping(BaseModel):
@@ -36,15 +36,15 @@ class ContainerMapping(BaseModel):
 
     python_node: str = Field(..., description="Python node class name")
     container_image: str = Field(..., description="Docker container image")
-    command: List[str] = Field(default_factory=list, description="Container command")
-    args: List[str] = Field(default_factory=list, description="Container arguments")
-    env: Dict[str, str] = Field(
+    command: list[str] = Field(default_factory=list, description="Container command")
+    args: list[str] = Field(default_factory=list, description="Container arguments")
+    env: dict[str, str] = Field(
         default_factory=dict, description="Environment variables"
     )
     resources: ResourceSpec = Field(
         default_factory=ResourceSpec, description="Resource specs"
     )
-    mount_paths: Dict[str, str] = Field(
+    mount_paths: dict[str, str] = Field(
         default_factory=dict, description="Volume mount paths"
     )
 
@@ -58,7 +58,7 @@ class ExportConfig(BaseModel):
     include_resources: bool = Field(True, description="Include resource specifications")
     validate_output: bool = Field(True, description="Validate exported format")
     container_registry: str = Field("", description="Container registry URL")
-    partial_export: Set[str] = Field(default_factory=set, description="Nodes to export")
+    partial_export: set[str] = Field(default_factory=set, description="Nodes to export")
 
 
 class NodeMapper:
@@ -71,7 +71,7 @@ class NodeMapper:
             ConfigurationException: If initialization fails
         """
         try:
-            self.mappings: Dict[str, ContainerMapping] = {}
+            self.mappings: dict[str, ContainerMapping] = {}
             self._initialize_default_mappings()
         except Exception as e:
             raise ConfigurationException(
@@ -198,7 +198,7 @@ class ExportValidator:
     """Validates exported workflow formats."""
 
     @staticmethod
-    def validate_yaml(data: Dict[str, Any]) -> bool:
+    def validate_yaml(data: dict[str, Any]) -> bool:
         """Validate YAML export format.
 
         Args:
@@ -271,7 +271,7 @@ class ExportValidator:
         return True
 
     @staticmethod
-    def validate_json(data: Dict[str, Any]) -> bool:
+    def validate_json(data: dict[str, Any]) -> bool:
         """Validate JSON export format.
 
         Args:
@@ -300,7 +300,7 @@ class ManifestGenerator:
 
     def generate_manifest(
         self, workflow: Workflow, node_mapper: NodeMapper
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Generate deployment manifest for a workflow.
 
         Args:
@@ -373,7 +373,7 @@ class ManifestGenerator:
 
     def _generate_node_spec(
         self, node_id: str, node_instance, node: Node, node_mapper: NodeMapper
-    ) -> Dict[str, Any]:
+    ) -> dict[str, Any]:
         """Generate node specification for manifest.
 
         Args:
@@ -479,7 +479,7 @@ class ManifestGenerator:
 class WorkflowExporter:
     """Main exporter for Kailash workflows."""
 
-    def __init__(self, config: Optional[ExportConfig] = None):
+    def __init__(self, config: ExportConfig | None = None):
         """Initialize the workflow exporter.
 
         Args:
@@ -506,7 +506,7 @@ class WorkflowExporter:
                 f"Failed to initialize workflow exporter: {e}"
             ) from e
 
-    def to_yaml(self, workflow: Workflow, output_path: Optional[str] = None) -> str:
+    def to_yaml(self, workflow: Workflow, output_path: str | None = None) -> str:
         """Export workflow to YAML format.
 
         Args:
@@ -552,7 +552,7 @@ class WorkflowExporter:
         except Exception as e:
             raise ExportException(f"Failed to export workflow to YAML: {e}") from e
 
-    def to_json(self, workflow: Workflow, output_path: Optional[str] = None) -> str:
+    def to_json(self, workflow: Workflow, output_path: str | None = None) -> str:
         """Export workflow to JSON format.
 
         Args:
@@ -598,7 +598,7 @@ class WorkflowExporter:
         except Exception as e:
             raise ExportException(f"Failed to export workflow to JSON: {e}") from e
 
-    def to_manifest(self, workflow: Workflow, output_path: Optional[str] = None) -> str:
+    def to_manifest(self, workflow: Workflow, output_path: str | None = None) -> str:
         """Export workflow as deployment manifest.
 
         Args:
@@ -645,7 +645,7 @@ class WorkflowExporter:
 
     def export_with_templates(
         self, workflow: Workflow, template_name: str, output_dir: str
-    ) -> Dict[str, str]:
+    ) -> dict[str, str]:
         """Export workflow using predefined templates.
 
         Args:
@@ -722,7 +722,7 @@ class WorkflowExporter:
 
         return exports
 
-    def _prepare_export_data(self, workflow: Workflow) -> Dict[str, Any]:
+    def _prepare_export_data(self, workflow: Workflow) -> dict[str, Any]:
         """Prepare workflow data for export.
 
         Args:
@@ -876,7 +876,7 @@ class WorkflowExporter:
 def export_workflow(
     workflow: Workflow,
     format: str = "yaml",
-    output_path: Optional[str] = None,
+    output_path: str | None = None,
     **config,
 ) -> str:
     """Export a workflow to specified format.

kailash/utils/templates.py

@@ -1,7 +1,6 @@
 """Project template system for Kailash SDK."""
 
 from pathlib import Path
-from typing import Dict, Optional
 
 from kailash.sdk_exceptions import TemplateError
 
@@ -266,7 +265,7 @@ __pycache__/
             },
         }
 
-    def get_template(self, template_name: str) -> Dict:
+    def get_template(self, template_name: str) -> dict:
         """Get an export template by name.
 
         Args:
@@ -286,7 +285,7 @@ __pycache__/
         self,
         project_name: str,
         template: str = "basic",
-        target_dir: Optional[str] = None,
+        target_dir: str | None = None,
     ) -> None:
         """Create a new project from a template.