ddeutil-workflow 0.0.37__py3-none-any.whl → 0.0.39__py3-none-any.whl

ddeutil/workflow/stages.py

@@ -23,6 +23,7 @@ template searching.
 """
 from __future__ import annotations
 
+import asyncio
 import contextlib
 import inspect
 import subprocess
@@ -31,11 +32,16 @@ import time
 import uuid
 from abc import ABC, abstractmethod
 from collections.abc import Iterator
+from concurrent.futures import (
+    Future,
+    ThreadPoolExecutor,
+    as_completed,
+)
 from inspect import Parameter
 from pathlib import Path
 from subprocess import CompletedProcess
 from textwrap import dedent
-from typing import Optional, Union
+from typing import Annotated, Optional, Union
 
 from pydantic import BaseModel, Field
 from pydantic.functional_validators import model_validator
@@ -43,25 +49,23 @@ from typing_extensions import Self
 
 from .__types import DictData, DictStr, TupleStr
 from .caller import TagFunc, extract_call
-from .conf import config, get_logger
-from .exceptions import StageException
+from .conf import config
+from .exceptions import StageException, to_dict
 from .result import Result, Status
 from .templates import not_in_template, param2template
 from .utils import (
-    cut_id,
     gen_id,
     make_exec,
 )
 
-logger = get_logger("ddeutil.workflow")
-
-
 __all__: TupleStr = (
     "EmptyStage",
     "BashStage",
     "PyStage",
     "CallStage",
     "TriggerStage",
+    "ForEachStage",
+    "ParallelStage",
     "Stage",
 )
 
@@ -127,13 +131,14 @@ class BaseStage(BaseModel, ABC):
         """Execute abstraction method that action something by sub-model class.
         This is important method that make this class is able to be the stage.
 
-        :param params: A parameter data that want to use in this execution.
+        :param params: (DictData) A parameter data that want to use in this
+            execution.
         :param result: (Result) A result object for keeping context and status
             data.
 
         :rtype: Result
         """
-        raise NotImplementedError("Stage should implement ``execute`` method.")
+        raise NotImplementedError("Stage should implement `execute` method.")
 
     def handler_execute(
         self,
@@ -142,8 +147,10 @@ class BaseStage(BaseModel, ABC):
         run_id: str | None = None,
         parent_run_id: str | None = None,
         result: Result | None = None,
+        raise_error: bool = False,
+        to: DictData | None = None,
     ) -> Result:
-        """Handler execution result from the stage `execute` method.
+        """Handler stage execution result from the stage `execute` method.
 
             This stage exception handler still use ok-error concept, but it
         allows you force catching an output result with error message by
@@ -151,26 +158,31 @@ class BaseStage(BaseModel, ABC):
 
             Execution   --> Ok      --> Result
                                         |-status: Status.SUCCESS
-                                        |-context:
-                                            |-outputs: ...
+                                        ╰-context:
+                                            ╰-outputs: ...
 
                         --> Error   --> Result (if env var was set)
                                         |-status: Status.FAILED
-                                        |-errors:
+                                        ╰-errors:
                                             |-class: ...
                                             |-name: ...
-                                            |-message: ...
+                                            ╰-message: ...
 
                         --> Error   --> Raise StageException(...)
 
             On the last step, it will set the running ID on a return result object
         from current stage ID before release the final result.
 
-        :param params: A parameter data that want to use in this execution.
+        :param params: (DictData) A parameterize value data that use in this
+            stage execution.
         :param run_id: (str) A running stage ID for this execution.
-        :param parent_run_id: A parent workflow running ID for this release.
+        :param parent_run_id: (str) A parent workflow running ID for this
+            execution.
         :param result: (Result) A result object for keeping context and status
-            data.
+            data before execution.
+        :param raise_error: (bool) A flag that all this method raise error
+        :param to: (DictData) A target object for auto set the return output
+            after execution.
 
         :rtype: Result
         """
@@ -178,18 +190,18 @@ class BaseStage(BaseModel, ABC):
             result,
             run_id=run_id,
             parent_run_id=parent_run_id,
-            id_logic=(self.name + (self.id or "")),
+            id_logic=self.iden,
         )
 
         try:
-            return self.execute(params, result=result)
+            rs: Result = self.execute(params, result=result)
+            if to is not None:
+                return self.set_outputs(rs.context, to=to)
+            return rs
         except Exception as err:
             result.trace.error(f"[STAGE]: {err.__class__.__name__}: {err}")
 
-            if config.stage_raise_error:
-                # NOTE: If error that raise from stage execution course by
-                #   itself, it will return that error with previous
-                #   dependency.
+            if raise_error or config.stage_raise_error:
                 if isinstance(err, StageException):
                     raise
 
@@ -198,22 +210,15 @@ class BaseStage(BaseModel, ABC):
                     f"{err.__class__.__name__}: {err}"
                 ) from None
 
-            # NOTE: Catching exception error object to result with
-            #   error_message and error keys.
-            return result.catch(
-                status=Status.FAILED,
-                context={
-                    "errors": {
-                        "class": err,
-                        "name": err.__class__.__name__,
-                        "message": f"{err.__class__.__name__}: {err}",
-                    },
-                },
-            )
+            errors: DictData = {"errors": to_dict(err)}
+            if to is not None:
+                return self.set_outputs(errors, to=to)
+
+            return result.catch(status=Status.FAILED, context=errors)
 
     def set_outputs(self, output: DictData, to: DictData) -> DictData:
         """Set an outputs from execution process to the received context. The
-        result from execution will pass to value of ``outputs`` key.
+        result from execution will pass to value of `outputs` key.
 
             For example of setting output method, If you receive execute output
         and want to set on the `to` like;
@@ -221,30 +226,29 @@ class BaseStage(BaseModel, ABC):
             ... (i)   output: {'foo': bar}
             ... (ii)  to: {}
 
-        The result of the `to` variable will be;
+            The result of the `to` argument will be;
 
             ... (iii) to: {
                         'stages': {
-                            '<stage-id>': {'outputs': {'foo': 'bar'}}
+                            '<stage-id>': {
+                                'outputs': {'foo': 'bar'},
+                                'skipped': False
+                            }
                         }
                     }
 
-        :param output: An output data that want to extract to an output key.
-        :param to: A context data that want to add output result.
+        :param output: (DictData) An output data that want to extract to an
+            output key.
+        :param to: (DictData) A context data that want to add output result.
+
         :rtype: DictData
         """
-        if self.id is None and not config.stage_default_id:
-            logger.warning(
-                "Output does not set because this stage does not set ID or "
-                "default stage ID config flag not be True."
-            )
-            return to
-
-        # NOTE: Create stages key to receive an output from the stage execution.
         if "stages" not in to:
             to["stages"] = {}
 
-        # NOTE: If the stage ID did not set, it will use its name instead.
+        if self.id is None and not config.stage_default_id:
+            return to
+
         _id: str = (
             param2template(self.id, params=to)
             if self.id
@@ -254,9 +258,12 @@ class BaseStage(BaseModel, ABC):
         errors: DictData = (
             {"errors": output.pop("errors", {})} if "errors" in output else {}
         )
-
-        # NOTE: Set the output to that stage generated ID with ``outputs`` key.
-        to["stages"][_id] = {"outputs": output, **errors}
+        skipping: dict[str, bool] = (
+            {"skipped": output.pop("skipped", False)}
+            if "skipped" in output
+            else {}
+        )
+        to["stages"][_id] = {"outputs": output, **skipping, **errors}
         return to
 
     def is_skipped(self, params: DictData | None = None) -> bool:
@@ -268,10 +275,11 @@ class BaseStage(BaseModel, ABC):
         :raise StageException: When return type of the eval condition statement
             does not return with boolean type.
 
-        :param params: A parameters that want to pass to condition template.
+        :param params: (DictData) A parameters that want to pass to condition
+            template.
+
         :rtype: bool
         """
-        # NOTE: Return false result if condition does not set.
         if self.condition is None:
             return False
 
@@ -299,6 +307,7 @@ class EmptyStage(BaseStage):
         >>> stage = {
         ...     "name": "Empty stage execution",
         ...     "echo": "Hello World",
+        ...     "sleep": 1,
         ... }
     """
 
@@ -308,7 +317,7 @@ class EmptyStage(BaseStage):
     )
     sleep: float = Field(
         default=0,
-        description="A second value to sleep before finish execution",
+        description="A second value to sleep before start execution",
         ge=0,
     )
 
@@ -322,13 +331,32 @@ class EmptyStage(BaseStage):
             The result context should be empty and do not process anything
         without calling logging function.
 
-        :param params: A context data that want to add output result. But this
-            stage does not pass any output.
+        :param params: (DictData) A context data that want to add output result.
+            But this stage does not pass any output.
         :param result: (Result) A result object for keeping context and status
             data.
 
         :rtype: Result
         """
+        result: Result = result or Result(
+            run_id=gen_id(self.name + (self.id or ""), unique=True)
+        )
+
+        result.trace.info(
+            f"[STAGE]: Empty-Execute: {self.name!r}: "
+            f"( {param2template(self.echo, params=params) or '...'} )"
+        )
+        if self.sleep > 0:
+            if self.sleep > 5:
+                result.trace.info(f"[STAGE]: ... sleep ({self.sleep} seconds)")
+            time.sleep(self.sleep)
+
+        return result.catch(status=Status.SUCCESS)
+
+    # TODO: Draft async execute method for the perf improvement.
+    async def aexecute(
+        self, params: DictData, *, result: Result | None = None
+    ) -> Result:  # pragma: no cov
         if result is None:  # pragma: no cov
             result: Result = Result(
                 run_id=gen_id(self.name + (self.id or ""), unique=True)
@@ -338,11 +366,8 @@ class EmptyStage(BaseStage):
             f"[STAGE]: Empty-Execute: {self.name!r}: "
             f"( {param2template(self.echo, params=params) or '...'} )"
         )
-        if self.sleep > 0:
-            if self.sleep > 30:
-                result.trace.info(f"[STAGE]: ... sleep ({self.sleep} seconds)")
-            time.sleep(self.sleep)
 
+        await asyncio.sleep(1)
         return result.catch(status=Status.SUCCESS)
 
 
@@ -382,20 +407,18 @@ class BashStage(BaseStage):
         step will write the `.sh` file before giving this file name to context.
         After that, it will auto delete this file automatic.
 
-        :param bash: A bash statement that want to execute.
-        :param env: An environment variable that use on this bash statement.
-        :param run_id: A running stage ID that use for writing sh file instead
-            generate by UUID4.
+        :param bash: (str) A bash statement that want to execute.
+        :param env: (DictStr) An environment variable that use on this bash
+            statement.
+        :param run_id: (str | None) A running stage ID that use for writing sh
+            file instead generate by UUID4.
+
         :rtype: Iterator[TupleStr]
         """
         run_id: str = run_id or uuid.uuid4()
         f_name: str = f"{run_id}.sh"
         f_shebang: str = "bash" if sys.platform.startswith("win") else "sh"
 
-        logger.debug(
-            f"({cut_id(run_id)}) [STAGE]: Start create `{f_name}` file."
-        )
-
         with open(f"./{f_name}", mode="w", newline="\n") as f:
             # NOTE: write header of `.sh` file
             f.write(f"#!/bin/{f_shebang}\n\n")
@@ -439,9 +462,11 @@ class BashStage(BaseStage):
             env=param2template(self.env, params),
             run_id=result.run_id,
         ) as sh:
+            result.trace.debug(f"... Start create `{sh[1]}` file.")
             rs: CompletedProcess = subprocess.run(
                 sh, shell=False, capture_output=True, text=True
             )
+
         if rs.returncode > 0:
             # NOTE: Prepare stderr message that returning from subprocess.
             err: str = (
@@ -457,8 +482,8 @@ class BashStage(BaseStage):
             status=Status.SUCCESS,
             context={
                 "return_code": rs.returncode,
-                "stdout": rs.stdout.rstrip("\n") or None,
-                "stderr": rs.stderr.rstrip("\n") or None,
+                "stdout": None if (out := rs.stdout.strip("\n")) == "" else out,
+                "stderr": None if (out := rs.stderr.strip("\n")) == "" else out,
             },
         )
 
@@ -515,24 +540,17 @@ class PyStage(BaseStage):
         """Override set an outputs method for the Python execution process that
         extract output from all the locals values.
 
-        :param output: An output data that want to extract to an output key.
-        :param to: A context data that want to add output result.
+        :param output: (DictData) An output data that want to extract to an
+            output key.
+        :param to: (DictData) A context data that want to add output result.
 
         :rtype: DictData
         """
-        # NOTE: The output will fileter unnecessary keys from locals.
-        lc: DictData = output.get("locals", {})
+        lc: DictData = output.pop("locals", {})
+        gb: DictData = output.pop("globals", {})
         super().set_outputs(
-            (
-                {k: lc[k] for k in self.filter_locals(lc)}
-                | {k: output[k] for k in output if k.startswith("error")}
-            ),
-            to=to,
+            {k: lc[k] for k in self.filter_locals(lc)} | output, to=to
         )
-
-        # NOTE: Override value that changing from the globals that pass via the
-        #   exec function.
-        gb: DictData = output.get("globals", {})
         to.update({k: gb[k] for k in to if k in gb})
         return to
 
@@ -553,24 +571,27 @@ class PyStage(BaseStage):
                 run_id=gen_id(self.name + (self.id or ""), unique=True)
             )
 
-        # NOTE: Replace the run statement that has templating value.
-        run: str = param2template(dedent(self.run), params)
-
-        # NOTE: create custom globals value that will pass to exec function.
-        _globals: DictData = (
-            globals() | params | param2template(self.vars, params)
-        )
         lc: DictData = {}
+        gb: DictData = (
+            globals()
+            | params
+            | param2template(self.vars, params)
+            | {"result": result}
+        )
 
         # NOTE: Start exec the run statement.
         result.trace.info(f"[STAGE]: Py-Execute: {self.name}")
+        result.trace.warning(
+            "[STAGE]: This stage allow use `eval` function, so, please "
+            "check your statement be safe before execute."
+        )
 
         # WARNING: The exec build-in function is very dangerous. So, it
         #   should use the re module to validate exec-string before running.
-        exec(run, _globals, lc)
+        exec(param2template(dedent(self.run), params), gb, lc)
 
         return result.catch(
-            status=Status.SUCCESS, context={"locals": lc, "globals": _globals}
+            status=Status.SUCCESS, context={"locals": lc, "globals": gb}
         )
 
 
@@ -583,11 +604,16 @@ class CallStage(BaseStage):
     statement. So, you can create your function complexly that you can for your
     objective to invoked by this stage object.
 
+        This stage is the usefull stage for run every job by a custom requirement
+    that you want by creating the Python function and adding it to the task
+    registry by importer syntax like `module.tasks.registry` not path style like
+    `module/tasks/registry`.
+
     Data Validate:
         >>> stage = {
         ...     "name": "Task stage execution",
         ...     "uses": "tasks/function-name@tag-name",
-        ...     "args": {"FOO": "BAR"},
+        ...     "args": {"arg01": "BAR", "kwarg01": 10},
         ... }
     """
 
@@ -631,15 +657,26 @@ class CallStage(BaseStage):
         #   calling.
         args: DictData = {"result": result} | param2template(self.args, params)
         ips = inspect.signature(t_func)
+        necessary_params: list[str] = [
+            k
+            for k in ips.parameters
+            if (
+                (v := ips.parameters[k]).default == Parameter.empty
+                and (
+                    v.kind != Parameter.VAR_KEYWORD
+                    or v.kind != Parameter.VAR_POSITIONAL
+                )
+            )
+        ]
         if any(
             (k.removeprefix("_") not in args and k not in args)
-            for k in ips.parameters
-            if ips.parameters[k].default == Parameter.empty
+            for k in necessary_params
         ):
             raise ValueError(
-                f"Necessary params, ({', '.join(ips.parameters.keys())}, ), "
+                f"Necessary params, ({', '.join(necessary_params)}, ), "
                 f"does not set to args"
             )
+
         # NOTE: add '_' prefix if it wants to use.
         for k in ips.parameters:
             if k.removeprefix("_") in args:
@@ -649,7 +686,13 @@ class CallStage(BaseStage):
             args.pop("result")
 
         result.trace.info(f"[STAGE]: Call-Execute: {t_func.name}@{t_func.tag}")
-        rs: DictData = t_func(**param2template(args, params))
+        if inspect.iscoroutinefunction(t_func):  # pragma: no cov
+            loop = asyncio.get_event_loop()
+            rs: DictData = loop.run_until_complete(
+                t_func(**param2template(args, params))
+            )
+        else:
+            rs: DictData = t_func(**param2template(args, params))
 
         # VALIDATE:
         #   Check the result type from call function, it should be dict.
@@ -717,54 +760,136 @@ class TriggerStage(BaseStage):
         )
 
 
-# NOTE:
-#   An order of parsing stage model on the Job model with ``stages`` field.
-#   From the current build-in stages, they do not have stage that have the same
-#   fields that because of parsing on the Job's stages key.
-#
-Stage = Union[
-    PyStage,
-    BashStage,
-    CallStage,
-    TriggerStage,
-    EmptyStage,
-]
-
-
-# TODO: Not implement this stages yet
 class ParallelStage(BaseStage):  # pragma: no cov
     """Parallel execution stage that execute child stages with parallel.
 
+        This stage is not the low-level stage model because it runs muti-stages
+    in this stage execution.
+
     Data Validate:
         >>> stage = {
         ...     "name": "Parallel stage execution.",
-        ...     "parallel": [
-        ...         {
-        ...             "name": "Echo first stage",
-        ...             "echo": "Start run with branch 1",
-        ...             "sleep": 3,
-        ...         },
-        ...         {
-        ...             "name": "Echo second stage",
-        ...             "echo": "Start run with branch 2",
-        ...             "sleep": 1,
-        ...         },
-        ...     ]
+        ...     "parallel": {
+        ...         "branch01": [
+        ...             {
+        ...                 "name": "Echo first stage",
+        ...                 "echo": "Start run with branch 1",
+        ...                 "sleep": 3,
+        ...             },
+        ...         ],
+        ...         "branch02": [
+        ...             {
+        ...                 "name": "Echo second stage",
+        ...                 "echo": "Start run with branch 2",
+        ...                 "sleep": 1,
+        ...             },
+        ...         ],
+        ...     }
         ... }
     """
 
-    parallel: list[Stage]
+    parallel: dict[str, list[Stage]] = Field(
+        description="A mapping of parallel branch ID.",
+    )
     max_parallel_core: int = Field(default=2)
 
+    @staticmethod
+    def task(
+        branch: str,
+        params: DictData,
+        result: Result,
+        stages: list[Stage],
+    ) -> DictData:
+        """Task execution method for passing a branch to each thread.
+
+        :param branch: A branch ID.
+        :param params: A parameter data that want to use in this execution.
+        :param result: (Result) A result object for keeping context and status
+            data.
+        :param stages:
+
+        :rtype: DictData
+        """
+        context = {"branch": branch, "stages": {}}
+        result.trace.debug(f"[STAGE]: Execute parallel branch: {branch!r}")
+        for stage in stages:
+            try:
+                stage.set_outputs(
+                    stage.handler_execute(
+                        params=params,
+                        run_id=result.run_id,
+                        parent_run_id=result.parent_run_id,
+                    ).context,
+                    to=context,
+                )
+            except StageException as err:  # pragma: no cov
+                result.trace.error(
+                    f"[STAGE]: Catch:\n\t{err.__class__.__name__}:" f"\n\t{err}"
+                )
+                context.update(
+                    {
+                        "errors": {
+                            "class": err,
+                            "name": err.__class__.__name__,
+                            "message": f"{err.__class__.__name__}: {err}",
+                        },
+                    },
+                )
+        return context
+
     def execute(
         self, params: DictData, *, result: Result | None = None
-    ) -> Result: ...
+    ) -> Result:
+        """Execute the stages that parallel each branch via multi-threading mode
+        or async mode by changing `async_mode` flag.
 
+        :param params: A parameter that want to pass before run any statement.
+        :param result: (Result) A result object for keeping context and status
+            data.
 
-# TODO: Not implement this stages yet
-class ForEachStage(BaseStage):  # pragma: no cov
-    """For-Each execution stage that execute child stages with an item in list of
-    item values.
+        :rtype: Result
+        """
+        if result is None:  # pragma: no cov
+            result: Result = Result(
+                run_id=gen_id(self.name + (self.id or ""), unique=True)
+            )
+
+        rs: DictData = {"parallel": {}}
+        status = Status.SUCCESS
+        with ThreadPoolExecutor(
+            max_workers=self.max_parallel_core,
+            thread_name_prefix="parallel_stage_exec_",
+        ) as executor:
+
+            futures: list[Future] = []
+            for branch in self.parallel:
+                futures.append(
+                    executor.submit(
+                        self.task,
+                        branch=branch,
+                        params=params,
+                        result=result,
+                        stages=self.parallel[branch],
+                    )
+                )
+
+            done = as_completed(futures, timeout=1800)
+            for future in done:
+                context: DictData = future.result()
+                rs["parallel"][context.pop("branch")] = context
+
+                if "errors" in context:
+                    status = Status.FAILED
+
+        return result.catch(status=status, context=rs)
+
+
+class ForEachStage(BaseStage):
+    """For-Each execution stage that execute child stages with an item in list
+    of item values.
+
+        This stage is not the low-level stage model because it runs muti-stages
+    in this stage execution.
 
     Data Validate:
         >>> stage = {
@@ -779,14 +904,126 @@ class ForEachStage(BaseStage):  # pragma: no cov
         ... }
     """
 
-    foreach: list[str]
-    stages: list[Stage]
+    foreach: Union[list[str], list[int]] = Field(
+        description=(
+            "A items for passing to each stages via ${{ item }} template."
+        ),
+    )
+    stages: list[Stage] = Field(
+        description=(
+            "A list of stage that will run with each item in the foreach field."
+        ),
+    )
+
+    def execute(
+        self, params: DictData, *, result: Result | None = None
+    ) -> Result:
+        """Execute the stages that pass each item form the foreach field.
+
+        :param params: A parameter that want to pass before run any statement.
+        :param result: (Result) A result object for keeping context and status
+            data.
+
+        :rtype: Result
+        """
+        if result is None:  # pragma: no cov
+            result: Result = Result(
+                run_id=gen_id(self.name + (self.id or ""), unique=True)
+            )
+
+        rs: DictData = {"items": self.foreach, "foreach": {}}
+        status = Status.SUCCESS
+        for item in self.foreach:
+            result.trace.debug(f"[STAGE]: Execute foreach item: {item!r}")
+            params["item"] = item
+            context = {"stages": {}}
+
+            for stage in self.stages:
+                try:
+                    stage.set_outputs(
+                        stage.handler_execute(
+                            params=params,
+                            run_id=result.run_id,
+                            parent_run_id=result.parent_run_id,
+                        ).context,
+                        to=context,
+                    )
+                except StageException as err:  # pragma: no cov
+                    status = Status.FAILED
+                    result.trace.error(
+                        f"[STAGE]: Catch:\n\t{err.__class__.__name__}:"
+                        f"\n\t{err}"
+                    )
+                    context.update(
+                        {
+                            "errors": {
+                                "class": err,
+                                "name": err.__class__.__name__,
+                                "message": f"{err.__class__.__name__}: {err}",
+                            },
+                        },
+                    )
+
+            rs["foreach"][item] = context
+
+        return result.catch(status=status, context=rs)
+
+
+# TODO: Not implement this stages yet
+class IfStage(BaseStage):  # pragma: no cov
+    """If execution stage.
+
+    Data Validate:
+        >>> stage = {
+        ...     "name": "If stage execution.",
+        ...     "case": "${{ param.test }}",
+        ...     "match": [
+        ...         {
+        ...             "case": "1",
+        ...             "stage": {
+        ...                 "name": "Stage case 1",
+        ...                 "eche": "Hello case 1",
+        ...             },
+        ...         },
+        ...         {
+        ...             "case": "2",
+        ...             "stage": {
+        ...                 "name": "Stage case 2",
+        ...                 "eche": "Hello case 2",
+        ...             },
+        ...         },
+        ...         {
+        ...             "case": "_",
+        ...             "stage": {
+        ...                 "name": "Stage else",
+        ...                 "eche": "Hello case else",
+        ...             },
+        ...         },
+        ...     ],
+        ... }
+
+    """
+
+    case: str = Field(description="A case condition for routing.")
+    match: list[dict[str, Union[str, Stage]]]
 
     def execute(
         self, params: DictData, *, result: Result | None = None
     ) -> Result: ...
 
 
+class RaiseStage(BaseStage):  # pragma: no cov
+    message: str = Field(
+        description="An error message that want to raise",
+        alias="raise",
+    )
+
+    def execute(
+        self, params: DictData, *, result: Result | None = None
+    ) -> Result:
+        raise StageException(self.message)
+
+
 # TODO: Not implement this stages yet
 class HookStage(BaseStage):  # pragma: no cov
     hook: str
@@ -820,3 +1057,36 @@ class VirtualPyStage(PyStage):  # pragma: no cov
     vars: DictData
 
     def create_py_file(self, py: str, run_id: str | None): ...
+
+    def execute(
+        self, params: DictData, *, result: Result | None = None
+    ) -> Result:
+        return super().execute(params, result=result)
+
+
+# TODO: Not implement this stages yet
+class SensorStage(BaseStage):  # pragma: no cov
+
+    def execute(
+        self, params: DictData, *, result: Result | None = None
+    ) -> Result: ...
+
+
+# NOTE:
+#   An order of parsing stage model on the Job model with ``stages`` field.
+#   From the current build-in stages, they do not have stage that have the same
+#   fields that because of parsing on the Job's stages key.
+#
+Stage = Annotated[
+    Union[
+        EmptyStage,
+        BashStage,
+        CallStage,
+        TriggerStage,
+        ForEachStage,
+        ParallelStage,
+        PyStage,
+        RaiseStage,
+    ],
+    Field(union_mode="smart"),
+]