stabilize 0.9.2__py3-none-any.whl

stabilize/tasks/interface.py

@@ -0,0 +1,335 @@
+"""
+Task interface definitions.
+
+This module defines the Task interface and its variants (RetryableTask,
+SkippableTask) that all task implementations must follow.
+"""
+
+from __future__ import annotations
+
+from abc import ABC, abstractmethod
+from collections.abc import Callable
+from datetime import timedelta
+from typing import TYPE_CHECKING
+
+from stabilize.tasks.result import TaskResult
+
+if TYPE_CHECKING:
+    from stabilize.models.stage import StageExecution
+
+
+class Task(ABC):
+    """
+    Base interface for all tasks.
+
+    Tasks are the atomic units of work in a pipeline. Each task:
+    - Receives the current stage context
+    - Performs some work
+    - Returns a TaskResult indicating status and any outputs
+
+    Example:
+        class DeployTask(Task):
+            def execute(self, stage: StageExecution) -> TaskResult:
+                # Get inputs from context
+                cluster = stage.context.get("cluster")
+                image = stage.context.get("image")
+
+                # Do the work
+                deployment_id = deploy(cluster, image)
+
+                # Return result with outputs
+                return TaskResult.success(
+                    outputs={"deploymentId": deployment_id}
+                )
+    """
+
+    @abstractmethod
+    def execute(self, stage: StageExecution) -> TaskResult:
+        """
+        Execute the task.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            TaskResult indicating status and any outputs
+
+        Raises:
+            Exception: Any exception will be caught and handled by the runner
+        """
+        pass
+
+    def on_timeout(self, stage: StageExecution) -> TaskResult | None:
+        """
+        Called when the task times out.
+
+        Override to provide custom timeout handling. If None is returned,
+        the default timeout behavior applies.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            Optional TaskResult to use instead of default timeout
+        """
+        return None
+
+    def on_cancel(self, stage: StageExecution) -> TaskResult | None:
+        """
+        Called when the execution is canceled.
+
+        Override to provide cleanup logic when execution is canceled.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            Optional TaskResult with cleanup results
+        """
+        return None
+
+    @property
+    def aliases(self) -> list[str]:
+        """
+        Alternative names for this task type.
+
+        Used for backward compatibility when task types are renamed.
+
+        Returns:
+            List of alternative type names
+        """
+        return []
+
+
+class RetryableTask(Task):
+    """
+    A task that can be retried with timeout and backoff.
+
+    Retryable tasks return RUNNING status while waiting for some condition.
+    They are re-executed after a backoff period until they succeed, fail,
+    or timeout.
+
+    Example:
+        class WaitForDeployTask(RetryableTask):
+            def get_timeout(self) -> timedelta:
+                return timedelta(minutes=30)
+
+            def get_backoff_period(self, stage, duration) -> timedelta:
+                return timedelta(seconds=10)
+
+            def execute(self, stage: StageExecution) -> TaskResult:
+                deployment_id = stage.context.get("deploymentId")
+                status = check_deployment_status(deployment_id)
+
+                if status == "complete":
+                    return TaskResult.success()
+                elif status == "failed":
+                    return TaskResult.terminal("Deployment failed")
+                else:
+                    return TaskResult.running()
+    """
+
+    @abstractmethod
+    def get_timeout(self) -> timedelta:
+        """
+        Get the maximum time this task can run before timing out.
+
+        Returns:
+            Maximum execution time
+        """
+        pass
+
+    def get_backoff_period(
+        self,
+        stage: StageExecution,
+        duration: timedelta,
+    ) -> timedelta:
+        """
+        Get the backoff period before retrying.
+
+        Override to implement dynamic backoff based on how long
+        the task has been running.
+
+        Args:
+            stage: The stage execution context
+            duration: How long the task has been running
+
+        Returns:
+            Time to wait before retrying
+        """
+        return timedelta(seconds=1)
+
+    def get_dynamic_timeout(self, stage: StageExecution) -> timedelta:
+        """
+        Get dynamic timeout based on stage context.
+
+        Override to implement context-based timeouts.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            Timeout duration
+        """
+        return self.get_timeout()
+
+    def get_dynamic_backoff_period(
+        self,
+        stage: StageExecution,
+        duration: timedelta,
+    ) -> timedelta:
+        """
+        Get dynamic backoff based on stage context.
+
+        Args:
+            stage: The stage execution context
+            duration: How long the task has been running
+
+        Returns:
+            Time to wait before retrying
+        """
+        return self.get_backoff_period(stage, duration)
+
+
+class OverridableTimeoutRetryableTask(RetryableTask):
+    """
+    A retryable task whose timeout can be overridden by the stage.
+
+    The stage can set a 'stageTimeoutMs' context value to override
+    the default timeout.
+    """
+
+    def get_dynamic_timeout(self, stage: StageExecution) -> timedelta:
+        """Get timeout, potentially overridden by stage context."""
+        if "stageTimeoutMs" in stage.context:
+            return timedelta(milliseconds=stage.context["stageTimeoutMs"])
+        return self.get_timeout()
+
+
+class SkippableTask(Task):
+    """
+    A task that can be conditionally skipped.
+
+    Override is_enabled() to control when the task should be skipped.
+    """
+
+    def is_enabled(self, stage: StageExecution) -> bool:
+        """
+        Check if this task is enabled.
+
+        Override to implement skip logic.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            True if task should execute, False to skip
+        """
+        return True
+
+    def execute(self, stage: StageExecution) -> TaskResult:
+        """Execute the task if enabled."""
+        if not self.is_enabled(stage):
+            return TaskResult.skipped()
+        return self.do_execute(stage)
+
+    @abstractmethod
+    def do_execute(self, stage: StageExecution) -> TaskResult:
+        """
+        Perform the actual task execution.
+
+        Args:
+            stage: The stage execution context
+
+        Returns:
+            TaskResult indicating status
+        """
+        pass
+
+
+class CallableTask(Task):
+    """
+    A task that wraps a callable function.
+
+    Allows using simple functions as tasks without creating a class.
+
+    Example:
+        def my_task(stage: StageExecution) -> TaskResult:
+            return TaskResult.success(outputs={"result": "done"})
+
+        task = CallableTask(my_task)
+    """
+
+    def __init__(
+        self,
+        func: Callable[[StageExecution], TaskResult],
+        name: str | None = None,
+    ) -> None:
+        """
+        Initialize with a callable.
+
+        Args:
+            func: The function to call
+            name: Optional name for the task
+        """
+        self._func = func
+        self._name = name or func.__name__
+
+    def execute(self, stage: StageExecution) -> TaskResult:
+        """Execute the wrapped function."""
+        return self._func(stage)
+
+    @property
+    def name(self) -> str:
+        """Get the task name."""
+        return self._name
+
+
+class NoOpTask(Task):
+    """
+    A task that does nothing.
+
+    Useful for testing or placeholder stages.
+    """
+
+    def execute(self, stage: StageExecution) -> TaskResult:
+        """Return success immediately."""
+        return TaskResult.success()
+
+
+class WaitTask(RetryableTask):
+    """
+    A task that waits for a specified duration.
+
+    Reads 'waitTime' from stage context (in seconds).
+    """
+
+    def get_timeout(self) -> timedelta:
+        """Wait tasks have a long timeout."""
+        return timedelta(hours=24)
+
+    def get_backoff_period(
+        self,
+        stage: StageExecution,
+        duration: timedelta,
+    ) -> timedelta:
+        """Check every second."""
+        return timedelta(seconds=1)
+
+    def execute(self, stage: StageExecution) -> TaskResult:
+        """Wait for the specified time."""
+        import time
+
+        wait_time = stage.context.get("waitTime", 0)
+        start_time = stage.context.get("waitStartTime")
+        current_time = int(time.time())
+
+        if start_time is None:
+            # First execution - record start time
+            return TaskResult.running(context={"waitStartTime": current_time})
+
+        elapsed = current_time - start_time
+        if elapsed >= wait_time:
+            return TaskResult.success(outputs={"waitedSeconds": elapsed})
+
+        return TaskResult.running()

stabilize/tasks/registry.py

@@ -0,0 +1,255 @@
+"""
+Task registry for resolving task implementations.
+
+This module provides the TaskRegistry class for registering and resolving
+task implementations by name or type.
+"""
+
+from __future__ import annotations
+
+import logging
+from collections.abc import Callable
+
+from stabilize.tasks.interface import CallableTask, Task
+from stabilize.tasks.result import TaskResult
+
+if False:  # TYPE_CHECKING
+    from stabilize.models.stage import StageExecution
+
+logger = logging.getLogger(__name__)
+
+# Type for task callable
+TaskCallable = Callable[["StageExecution"], TaskResult]
+
+# Type for task implementation - can be a Task class or callable
+TaskImplementation = type[Task] | Task | TaskCallable
+
+
+class TaskNotFoundError(Exception):
+    """Raised when a task type cannot be resolved."""
+
+    def __init__(self, task_type: str):
+        self.task_type = task_type
+        super().__init__(f"No task found for type: {task_type}")
+
+
+class TaskRegistry:
+    """
+    Registry for task implementations.
+
+    Allows registering tasks by name and resolving them at runtime.
+    Supports:
+    - Task classes (instantiated on resolve)
+    - Task instances (used directly)
+    - Callable functions (wrapped in CallableTask)
+
+    Example:
+        registry = TaskRegistry()
+
+        # Register a task class
+        registry.register("deploy", DeployTask)
+
+        # Register a task instance
+        registry.register("notify", NotifyTask(slack_client))
+
+        # Register a function
+        @registry.task("validate")
+        def validate_inputs(stage):
+            return TaskResult.success()
+
+        # Resolve and use
+        task = registry.get("deploy")
+        result = task.execute(stage)
+    """
+
+    def __init__(self) -> None:
+        self._tasks: dict[str, TaskImplementation] = {}
+        self._aliases: dict[str, str] = {}
+
+    def register(
+        self,
+        name: str,
+        task: TaskImplementation,
+        aliases: list[str] | None = None,
+    ) -> None:
+        """
+        Register a task implementation.
+
+        Args:
+            name: The task type name
+            task: Task class, instance, or callable
+            aliases: Optional alternative names
+
+        Raises:
+            ValueError: If name is already registered
+        """
+        if name in self._tasks:
+            logger.warning(f"Overwriting existing task registration: {name}")
+
+        self._tasks[name] = task
+
+        # Register aliases
+        if aliases:
+            for alias in aliases:
+                self._aliases[alias] = name
+
+        # Check for aliases on the task itself
+        if isinstance(task, type) and issubclass(task, Task):
+            instance = task()
+            for alias in instance.aliases:
+                self._aliases[alias] = name
+        elif isinstance(task, Task):
+            for alias in task.aliases:
+                self._aliases[alias] = name
+
+        logger.debug(f"Registered task: {name}")
+
+    def register_class(
+        self,
+        task_class: type[Task],
+        name: str | None = None,
+    ) -> None:
+        """
+        Register a task class using its class name.
+
+        Args:
+            task_class: The task class to register
+            name: Optional name override
+        """
+        task_name = name or task_class.__name__
+        self.register(task_name, task_class)
+
+    def task(
+        self,
+        name: str,
+        aliases: list[str] | None = None,
+    ) -> Callable[[TaskCallable], TaskCallable]:
+        """
+        Decorator to register a function as a task.
+
+        Args:
+            name: The task type name
+            aliases: Optional alternative names
+
+        Returns:
+            Decorator function
+
+        Example:
+            @registry.task("validate")
+            def validate_inputs(stage):
+                return TaskResult.success()
+        """
+
+        def decorator(func: TaskCallable) -> TaskCallable:
+            self.register(name, func, aliases)
+            return func
+
+        return decorator
+
+    def get(self, name: str) -> Task:
+        """
+        Get a task implementation by name.
+
+        Args:
+            name: The task type name
+
+        Returns:
+            A Task instance
+
+        Raises:
+            TaskNotFoundError: If task not found
+        """
+        # Check aliases first
+        resolved_name = self._aliases.get(name, name)
+
+        if resolved_name not in self._tasks:
+            raise TaskNotFoundError(resolved_name)
+
+        impl = self._tasks[resolved_name]
+
+        # Handle different registration types
+        if isinstance(impl, Task):
+            return impl
+        elif isinstance(impl, type) and issubclass(impl, Task):
+            return impl()
+        elif callable(impl):
+            # Cast to the proper type for CallableTask
+            from typing import cast
+
+            func = cast(TaskCallable, impl)
+            return CallableTask(func, name=resolved_name)
+        else:
+            raise TaskNotFoundError(resolved_name)
+
+    def get_by_class(self, class_name: str) -> Task:
+        """
+        Get a task by its implementing class name.
+
+        Args:
+            class_name: Fully qualified or simple class name
+
+        Returns:
+            A Task instance
+
+        Raises:
+            TaskNotFoundError: If task not found
+        """
+        # Try exact match first
+        if class_name in self._tasks:
+            return self.get(class_name)
+
+        # Try simple class name
+        simple_name = class_name.split(".")[-1]
+        if simple_name in self._tasks:
+            return self.get(simple_name)
+
+        # Try lowercase
+        if simple_name.lower() in self._tasks:
+            return self.get(simple_name.lower())
+
+        raise TaskNotFoundError(class_name)
+
+    def has(self, name: str) -> bool:
+        """Check if a task is registered."""
+        resolved_name = self._aliases.get(name, name)
+        return resolved_name in self._tasks
+
+    def list_tasks(self) -> list[str]:
+        """Get all registered task names."""
+        return list(self._tasks.keys())
+
+    def clear(self) -> None:
+        """Clear all registrations."""
+        self._tasks.clear()
+        self._aliases.clear()
+
+
+# Global registry instance
+_default_registry: TaskRegistry | None = None
+
+
+def get_default_registry() -> TaskRegistry:
+    """Get the default global task registry."""
+    global _default_registry
+    if _default_registry is None:
+        _default_registry = TaskRegistry()
+    return _default_registry
+
+
+def register_task(
+    name: str,
+    task: TaskImplementation,
+    aliases: list[str] | None = None,
+) -> None:
+    """Register a task in the default registry."""
+    get_default_registry().register(name, task, aliases)
+
+
+def get_task(name: str) -> Task:
+    """Get a task from the default registry."""
+    return get_default_registry().get(name)
+
+
+def task(name: str, aliases: list[str] | None = None) -> Callable[[TaskCallable], TaskCallable]:
+    """Decorator to register a task in the default registry."""
+    return get_default_registry().task(name, aliases)