ddeutil-workflow 0.0.77__py3-none-any.whl → 0.0.79__py3-none-any.whl

ddeutil/workflow/stages.py

@@ -77,7 +77,15 @@ from pathlib import Path
 from subprocess import CompletedProcess
 from textwrap import dedent
 from threading import Event
-from typing import Annotated, Any, Optional, TypeVar, Union, get_type_hints
+from typing import (
+    Annotated,
+    Any,
+    Callable,
+    Optional,
+    TypeVar,
+    Union,
+    get_type_hints,
+)
 
 from ddeutil.core import str2list
 from pydantic import BaseModel, Field, ValidationError
@@ -106,7 +114,7 @@ from .reusables import (
     not_in_template,
     param2template,
 )
-from .traces import Trace, get_trace
+from .traces import TraceManager, get_trace
 from .utils import (
     delay,
     dump_all,
@@ -177,7 +185,7 @@ class BaseStage(BaseModel, ABC):
             "A stage description that use to logging when start execution."
         ),
     )
-    condition: StrOrNone = Field(
+    condition: Optional[Union[str, bool]] = Field(
         default=None,
         description=(
             "A stage condition statement to allow stage executable. This field "
@@ -296,7 +304,7 @@ class BaseStage(BaseModel, ABC):
         parent_run_id: str = run_id
         run_id: str = run_id or gen_id(self.iden, unique=True)
         context: DictData = {"status": WAIT}
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         try:
@@ -513,6 +521,9 @@ class BaseStage(BaseModel, ABC):
         if not self.condition:
             return False
 
+        if isinstance(self.condition, bool):
+            return self.condition
+
         try:
             # WARNING: The eval build-in function is very dangerous. So, it
             #   should use the `re` module to validate eval-string before
@@ -623,7 +634,7 @@ class BaseAsyncStage(BaseStage, ABC):
         parent_run_id: StrOrNone = run_id
         run_id: str = run_id or gen_id(self.iden, unique=True)
         context: DictData = {}
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         try:
@@ -766,7 +777,7 @@ class BaseRetryStage(BaseAsyncStage, ABC):  # pragma: no cov
         current_retry: int = 0
         exception: Exception
         catch(context, status=WAIT)
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
 
@@ -839,7 +850,7 @@ class BaseRetryStage(BaseAsyncStage, ABC):  # pragma: no cov
         current_retry: int = 0
         exception: Exception
         catch(context, status=WAIT)
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
 
@@ -971,7 +982,7 @@ class EmptyStage(BaseAsyncStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         message: str = (
@@ -1024,7 +1035,7 @@ class EmptyStage(BaseAsyncStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         message: str = (
@@ -1189,7 +1200,7 @@ class BashStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         bash: str = param2template(
@@ -1253,7 +1264,7 @@ class BashStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         bash: str = param2template(
@@ -1394,7 +1405,7 @@ class PyStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.info("[STAGE]: Prepare `globals` and `locals` variables.")
@@ -1475,7 +1486,7 @@ class PyStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         await trace.ainfo("[STAGE]: Prepare `globals` and `locals` variables.")
@@ -1580,15 +1591,23 @@ class CallStage(BaseRetryStage):
 
         :rtype: Any
         """
-        if isinstance(value, dict):
-            if any(k in value for k in ("result", "extras")):
-                raise ValueError(
-                    "The argument on workflow template for the caller stage "
-                    "should not pass `result` and `extras`. They are special "
-                    "arguments."
-                )
+        if isinstance(value, dict) and any(
+            k in value for k in ("result", "extras")
+        ):
+            raise ValueError(
+                "The argument on workflow template for the caller stage "
+                "should not pass `result` and `extras`. They are special "
+                "arguments."
+            )
         return value
 
+    def get_caller(self, params: DictData) -> Callable[[], TagFunc]:
+        """Get the lazy TagFuc object from registry."""
+        return extract_call(
+            param2template(self.uses, params, extras=self.extras),
+            registries=self.extras.get("registry_caller"),
+        )
+
     def process(
         self,
         params: DictData,
@@ -1612,14 +1631,10 @@ class CallStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
-        call_func: TagFunc = extract_call(
-            param2template(self.uses, params, extras=self.extras),
-            registries=self.extras.get("registry_caller"),
-        )()
-
+        call_func: TagFunc = self.get_caller(params=params)()
         trace.info(f"[STAGE]: Caller Func: '{call_func.name}@{call_func.tag}'")
 
         # VALIDATE: check input task caller parameters that exists before
@@ -1677,7 +1692,7 @@ class CallStage(BaseRetryStage):
             )
 
         args: DictData = self.validate_model_args(
-            call_func, args, run_id, parent_run_id
+            call_func, args, run_id, parent_run_id, extras=self.extras
         )
         if inspect.iscoroutinefunction(call_func):
             loop = asyncio.get_event_loop()
@@ -1735,14 +1750,10 @@ class CallStage(BaseRetryStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
-        call_func: TagFunc = extract_call(
-            param2template(self.uses, params, extras=self.extras),
-            registries=self.extras.get("registry_caller"),
-        )()
-
+        call_func: TagFunc = self.get_caller(params=params)()
         await trace.ainfo(
             f"[STAGE]: Caller Func: '{call_func.name}@{call_func.tag}'"
         )
@@ -1795,8 +1806,13 @@ class CallStage(BaseRetryStage):
         if "extras" not in sig.parameters and not has_keyword:
             args.pop("extras")
 
+        if event and event.is_set():
+            raise StageCancelError(
+                "Execution was canceled from the event before start parallel."
+            )
+
         args: DictData = self.validate_model_args(
-            call_func, args, run_id, parent_run_id
+            call_func, args, run_id, parent_run_id, extras=self.extras
         )
         if inspect.iscoroutinefunction(call_func):
             rs: DictOrModel = await call_func(
@@ -1829,12 +1845,13 @@ class CallStage(BaseRetryStage):
             extras=self.extras,
         )
 
+    @staticmethod
     def validate_model_args(
-        self,
         func: TagFunc,
         args: DictData,
         run_id: str,
         parent_run_id: Optional[str] = None,
+        extras: Optional[DictData] = None,
     ) -> DictData:
         """Validate an input arguments before passing to the caller function.
 
@@ -1846,18 +1863,18 @@ class CallStage(BaseRetryStage):
         :rtype: DictData
         """
         try:
-            model_instance: BaseModel = create_model_from_caller(
-                func
-            ).model_validate(args)
-            override: DictData = dict(model_instance)
+            override: DictData = dict(
+                create_model_from_caller(func).model_validate(args)
+            )
             args.update(override)
+
             type_hints: dict[str, Any] = get_type_hints(func)
             for arg in type_hints:
 
                 if arg == "return":
                     continue
 
-                if arg.removeprefix("_") in args:
+                if arg.startswith("_") and arg.removeprefix("_") in args:
                     args[arg] = args.pop(arg.removeprefix("_"))
                     continue
 
@@ -1867,8 +1884,8 @@ class CallStage(BaseRetryStage):
                 "Validate argument from the caller function raise invalid type."
             ) from e
         except TypeError as e:
-            trace: Trace = get_trace(
-                run_id, parent_run_id=parent_run_id, extras=self.extras
+            trace: TraceManager = get_trace(
+                run_id, parent_run_id=parent_run_id, extras=extras
             )
             trace.warning(
                 f"[STAGE]: Get type hint raise TypeError: {e}, so, it skip "
@@ -1992,7 +2009,7 @@ class TriggerStage(BaseNestedStage):
         """
         from .workflow import Workflow
 
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         _trigger: str = param2template(self.trigger, params, extras=self.extras)
@@ -2097,7 +2114,7 @@ class ParallelStage(BaseNestedStage):
 
         :rtype: tuple[Status, DictData]
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.debug(f"[STAGE]: Execute Branch: {branch!r}")
@@ -2227,7 +2244,7 @@ class ParallelStage(BaseNestedStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         event: Event = event or Event()
@@ -2295,7 +2312,13 @@ class ForEachStage(BaseNestedStage):
         ... }
     """
 
-    foreach: Union[list[str], list[int], str] = Field(
+    foreach: Union[
+        list[str],
+        list[int],
+        str,
+        dict[str, Any],
+        dict[int, Any],
+    ] = Field(
         description=(
             "A items for passing to stages via ${{ item }} template parameter."
         ),
@@ -2360,7 +2383,7 @@ class ForEachStage(BaseNestedStage):
 
         :rtype: tuple[Status, Result]
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.debug(f"[STAGE]: Execute Item: {item!r}")
@@ -2497,11 +2520,11 @@ class ForEachStage(BaseNestedStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         event: Event = event or Event()
-        foreach: Union[list[str], list[int]] = pass_env(
+        foreach: Union[list[str], list[int], str] = pass_env(
             param2template(self.foreach, params, extras=self.extras)
         )
 
@@ -2516,9 +2539,10 @@ class ForEachStage(BaseNestedStage):
                 ) from e
 
         # [VALIDATE]: Type of the foreach should be `list` type.
-        elif not isinstance(foreach, list):
+        elif isinstance(foreach, dict):
             raise TypeError(
-                f"Does not support foreach: {foreach!r} ({type(foreach)})"
+                f"Does not support dict foreach: {foreach!r} ({type(foreach)}) "
+                f"yet."
             )
         # [Validate]: Value in the foreach item should not be duplicate when the
         #   `use_index_as_key` field did not set.
@@ -2681,7 +2705,7 @@ class UntilStage(BaseNestedStage):
         :rtype: tuple[Status, DictData, T]
         :return: Return a pair of Result and changed item.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.debug(f"[STAGE]: Execute Loop: {loop} (Item {item!r})")
@@ -2823,7 +2847,7 @@ class UntilStage(BaseNestedStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         event: Event = event or Event()
@@ -2975,7 +2999,7 @@ class CaseStage(BaseNestedStage):
 
         :rtype: DictData
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.debug(f"[STAGE]: Execute Case: {case!r}")
@@ -3056,7 +3080,7 @@ class CaseStage(BaseNestedStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
 
@@ -3154,7 +3178,7 @@ class RaiseStage(BaseAsyncStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         message: str = param2template(self.message, params, extras=self.extras)
@@ -3185,7 +3209,7 @@ class RaiseStage(BaseAsyncStage):
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         message: str = param2template(self.message, params, extras=self.extras)
@@ -3266,7 +3290,7 @@ class DockerStage(BaseStage):  # pragma: no cov
                 "by `pip install docker` first."
             ) from None
 
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         client = DockerClient(
@@ -3366,7 +3390,7 @@ class DockerStage(BaseStage):  # pragma: no cov
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         trace.info(f"[STAGE]: Docker: {self.image}:{self.tag}")
@@ -3476,7 +3500,7 @@ class VirtualPyStage(PyStage):  # pragma: no cov
         Returns:
             Result: The execution result with status and context data.
         """
-        trace: Trace = get_trace(
+        trace: TraceManager = get_trace(
             run_id, parent_run_id=parent_run_id, extras=self.extras
         )
         run: str = param2template(dedent(self.run), params, extras=self.extras)