ddeutil-workflow 0.0.8__py3-none-any.whl → 0.0.10__py3-none-any.whl

ddeutil/workflow/stage.py

@@ -19,7 +19,6 @@ from __future__ import annotations
 
 import contextlib
 import inspect
-import logging
 import os
 import subprocess
 import sys
@@ -31,14 +30,22 @@ from functools import wraps
 from inspect import Parameter
 from pathlib import Path
 from subprocess import CompletedProcess
+from textwrap import dedent
 from typing import Callable, Optional, Union
 
+try:
+    from typing import ParamSpec
+except ImportError:
+    from typing_extensions import ParamSpec
+
 from ddeutil.core import str2bool
 from pydantic import BaseModel, Field
 from pydantic.functional_validators import model_validator
+from typing_extensions import Self
 
 from .__types import DictData, DictStr, Re, TupleStr
 from .exceptions import StageException
+from .log import get_logger
 from .utils import (
     Registry,
     Result,
@@ -46,33 +53,66 @@ from .utils import (
     gen_id,
     make_exec,
     make_registry,
+    not_in_template,
     param2template,
 )
 
+P = ParamSpec("P")
+logger = get_logger("ddeutil.workflow")
+
+
+__all__: TupleStr = (
+    "Stage",
+    "EmptyStage",
+    "BashStage",
+    "PyStage",
+    "HookStage",
+    "TriggerStage",
+    "handler_result",
+)
+
 
-def handler_result(message: str | None = None):
-    """Decorator function for handler result from the stage execution."""
+def handler_result(message: str | None = None) -> Callable[P, Result]:
+    """Decorator function for handler result from the stage execution. This
+    function should to use with execution method only.
+
+    :param message: A message that want to add at prefix of exception statement.
+    """
     message: str = message or ""
 
-    def decorator(func):
+    def decorator(func: Callable[P, Result]) -> Callable[P, Result]:
 
         @wraps(func)
-        def wrapped(self: BaseStage, *args, **kwargs):
+        def wrapped(self: Stage, *args, **kwargs):
             try:
-                rs: DictData = func(self, *args, **kwargs)
-                return Result(status=0, context=rs)
+                # NOTE: Start calling origin function with a passing args.
+                return func(self, *args, **kwargs).set_run_id(self.run_id)
             except Exception as err:
-                logging.error(
+                # NOTE: Start catching error from the stage execution.
+                logger.error(
                     f"({self.run_id}) [STAGE]: {err.__class__.__name__}: {err}"
                 )
-                if isinstance(err, StageException):
+                if str2bool(
+                    os.getenv("WORKFLOW_CORE_STAGE_RAISE_ERROR", "true")
+                ):
+                    # NOTE: If error that raise from stage execution course by
+                    #   itself, it will return that error with previous
+                    #   dependency.
+                    if isinstance(err, StageException):
+                        raise StageException(
+                            f"{self.__class__.__name__}: {message}\n\t{err}"
+                        ) from err
                     raise StageException(
-                        f"{self.__class__.__name__}: {message}\n---\n\t{err}"
-                    ) from err
-                raise StageException(
-                    f"{self.__class__.__name__}: {message}\n---\n\t"
-                    f"{err.__class__.__name__}: {err}"
-                ) from None
+                        f"{self.__class__.__name__}: {message}\n\t"
+                        f"{err.__class__.__name__}: {err}"
+                    ) from None
+                rs: Result = Result(
+                    status=1,
+                    context={
+                        "error_message": f"{err.__class__.__name__}: {err}",
+                    },
+                )
+                return rs.set_run_id(self.run_id)
 
         return wrapped
 
@@ -93,24 +133,47 @@ class BaseStage(BaseModel, ABC):
         ),
     )
     name: str = Field(
-        description="A stage name that want to logging when start execution."
+        description="A stage name that want to logging when start execution.",
     )
     condition: Optional[str] = Field(
         default=None,
+        description="A stage condition statement to allow stage executable.",
         alias="if",
     )
     run_id: Optional[str] = Field(
         default=None,
         description="A running stage ID.",
         repr=False,
+        exclude=True,
     )
 
     @model_validator(mode="after")
     def __prepare_running_id(self):
+        """Prepare stage running ID that use default value of field and this
+        method will validate name and id fields should not contain any template
+        parameter (exclude matrix template).
+        """
         if self.run_id is None:
             self.run_id = gen_id(self.name + (self.id or ""), unique=True)
+
+        # VALIDATE: Validate stage id and name should not dynamic with params
+        #   template. (allow only matrix)
+        if not_in_template(self.id) or not_in_template(self.name):
+            raise ValueError(
+                "Stage name and ID should only template with matrix."
+            )
+
         return self
 
+    def get_running_id(self, run_id: str) -> Self:
+        """Return Stage model object that changing stage running ID with an
+        input running ID.
+
+        :param run_id: A replace stage running ID.
+        :rtype: Self
+        """
+        return self.model_copy(update={"run_id": run_id})
+
     @abstractmethod
     def execute(self, params: DictData) -> Result:
         """Execute abstraction method that action something by sub-model class.
@@ -121,41 +184,39 @@ class BaseStage(BaseModel, ABC):
         """
         raise NotImplementedError("Stage should implement ``execute`` method.")
 
-    def set_outputs(self, output: DictData, params: DictData) -> DictData:
+    def set_outputs(self, output: DictData, to: DictData) -> DictData:
         """Set an outputs from execution process to an input params.
 
         :param output: A output data that want to extract to an output key.
-        :param params: A context data that want to add output result.
+        :param to: A context data that want to add output result.
         :rtype: DictData
         """
         if not (
             self.id
-            or str2bool(os.getenv("WORKFLOW_CORE_DEFAULT_STAGE_ID", "false"))
+            or str2bool(os.getenv("WORKFLOW_CORE_STAGE_DEFAULT_ID", "false"))
         ):
-            logging.debug(
+            logger.debug(
                 f"({self.run_id}) [STAGE]: Output does not set because this "
                 f"stage does not set ID or default stage ID config flag not be "
                 f"True."
             )
-            return params
+            return to
 
         # NOTE: Create stages key to receive an output from the stage execution.
-        if "stages" not in params:
-            params["stages"] = {}
+        if "stages" not in to:
+            to["stages"] = {}
 
-        # TODO: Validate stage id and name should not dynamic with params
-        #   template. (allow only matrix)
         if self.id:
-            _id: str = param2template(self.id, params=params)
+            _id: str = param2template(self.id, params=to)
         else:
-            _id: str = gen_id(param2template(self.name, params=params))
+            _id: str = gen_id(param2template(self.name, params=to))
 
         # NOTE: Set the output to that stage generated ID.
-        params["stages"][_id] = {"outputs": output}
-        logging.debug(
+        logger.debug(
             f"({self.run_id}) [STAGE]: Set output complete with stage ID: {_id}"
         )
-        return params
+        to["stages"][_id] = {"outputs": output}
+        return to
 
     def is_skipped(self, params: DictData | None = None) -> bool:
         """Return true if condition of this stage do not correct.
@@ -174,8 +235,8 @@ class BaseStage(BaseModel, ABC):
                 raise TypeError("Return type of condition does not be boolean")
             return not rs
         except Exception as err:
-            logging.error(f"({self.run_id}) [STAGE]: {err}")
-            raise StageException(str(err)) from err
+            logger.error(f"({self.run_id}) [STAGE]: {err}")
+            raise StageException(f"{err.__class__.__name__}: {err}") from err
 
 
 class EmptyStage(BaseStage):
@@ -201,7 +262,7 @@ class EmptyStage(BaseStage):
         :param params: A context data that want to add output result. But this
             stage does not pass any output.
         """
-        logging.info(
+        logger.info(
             f"({self.run_id}) [STAGE]: Empty-Execute: {self.name!r}: "
             f"( {param2template(self.echo, params=params) or '...'} )"
         )
@@ -246,37 +307,40 @@ class BashStage(BaseStage):
         f_shebang: str = "bash" if sys.platform.startswith("win") else "sh"
         with open(f"./{f_name}", mode="w", newline="\n") as f:
             # NOTE: write header of `.sh` file
-            f.write(f"#!/bin/{f_shebang}\n")
+            f.write(f"#!/bin/{f_shebang}\n\n")
 
             # NOTE: add setting environment variable before bash skip statement.
             f.writelines([f"{k}='{env[k]}';\n" for k in env])
 
             # NOTE: make sure that shell script file does not have `\r` char.
-            f.write(bash.replace("\r\n", "\n"))
+            f.write("\n" + bash.replace("\r\n", "\n"))
 
+        # NOTE: Make this .sh file able to executable.
         make_exec(f"./{f_name}")
-        logging.debug(
+
+        logger.debug(
             f"({self.run_id}) [STAGE]: Start create `.sh` file and running a "
             f"bash statement."
         )
 
         yield [f_shebang, f_name]
 
+        # Note: Remove .sh file that use to run bash.
         Path(f"./{f_name}").unlink()
 
     @handler_result()
-    def execute(self, params: DictData) -> DictData:
+    def execute(self, params: DictData) -> Result:
         """Execute the Bash statement with the Python build-in ``subprocess``
         package.
 
         :param params: A parameter data that want to use in this execution.
         :rtype: Result
         """
-        bash: str = param2template(self.bash, params)
+        bash: str = param2template(dedent(self.bash), params)
         with self.__prepare_bash(
             bash=bash, env=param2template(self.env, params)
         ) as sh:
-            logging.info(f"({self.run_id}) [STAGE]: Shell-Execute: {sh}")
+            logger.info(f"({self.run_id}) [STAGE]: Shell-Execute: {sh}")
             rs: CompletedProcess = subprocess.run(
                 sh,
                 shell=False,
@@ -288,19 +352,19 @@ class BashStage(BaseStage):
                 rs.stderr.encode("utf-8").decode("utf-16")
                 if "\\x00" in rs.stderr
                 else rs.stderr
-            )
-            logging.error(
-                f"({self.run_id}) [STAGE]: {err}\n\n```bash\n{bash}```"
-            )
+            ).removesuffix("\n")
             raise StageException(
-                f"{err.__class__.__name__}: {err}\nRunning Statement:"
-                f"\n---\n```bash\n{bash}\n```"
+                f"Subprocess: {err}\nRunning Statement:\n---\n"
+                f"```bash\n{bash}\n```"
             )
-        return {
-            "return_code": rs.returncode,
-            "stdout": rs.stdout.rstrip("\n"),
-            "stderr": rs.stderr.rstrip("\n"),
-        }
+        return Result(
+            status=0,
+            context={
+                "return_code": rs.returncode,
+                "stdout": rs.stdout.rstrip("\n"),
+                "stderr": rs.stderr.rstrip("\n"),
+            },
+        )
 
 
 class PyStage(BaseStage):
@@ -327,48 +391,56 @@ class PyStage(BaseStage):
         ),
     )
 
-    def set_outputs(self, output: DictData, params: DictData) -> DictData:
+    def set_outputs(self, output: DictData, to: DictData) -> DictData:
         """Set an outputs from the Python execution process to an input params.
 
         :param output: A output data that want to extract to an output key.
-        :param params: A context data that want to add output result.
+        :param to: A context data that want to add output result.
         :rtype: DictData
         """
         # NOTE: The output will fileter unnecessary keys from locals.
         _locals: DictData = output["locals"]
         super().set_outputs(
-            {k: _locals[k] for k in _locals if k != "__annotations__"},
-            params=params,
+            {k: _locals[k] for k in _locals if k != "__annotations__"}, to=to
         )
 
         # NOTE:
         #   Override value that changing from the globals that pass via exec.
         _globals: DictData = output["globals"]
-        params.update({k: _globals[k] for k in params if k in _globals})
-        return params
+        to.update({k: _globals[k] for k in to if k in _globals})
+        return to
 
     @handler_result()
-    def execute(self, params: DictData) -> DictData:
+    def execute(self, params: DictData) -> Result:
         """Execute the Python statement that pass all globals and input params
         to globals argument on ``exec`` build-in function.
 
         :param params: A parameter that want to pass before run any statement.
         :rtype: Result
         """
+        # 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)
         )
         _locals: DictData = {}
-        run: str = param2template(self.run, params)
-        logging.info(f"({self.run_id}) [STAGE]: Py-Execute: {uuid.uuid4()}")
+
+        # NOTE: Start exec the run statement.
+        logger.info(f"({self.run_id}) [STAGE]: Py-Execute: {self.name}")
         exec(run, _globals, _locals)
-        return {"locals": _locals, "globals": _globals}
+
+        return Result(
+            status=0, context={"locals": _locals, "globals": _globals}
+        )
 
 
 @dataclass
 class HookSearch:
-    """Hook Search dataclass."""
+    """Hook Search dataclass that use for receive regular expression grouping
+    dict from searching hook string value.
+    """
 
     path: str
     func: str
@@ -376,13 +448,16 @@ class HookSearch:
 
 
 def extract_hook(hook: str) -> Callable[[], TagFunc]:
-    """Extract Hook string value to hook function.
+    """Extract Hook function from string value to hook partial function that
+    does run it at runtime.
 
     :param hook: A hook value that able to match with Task regex.
     :rtype: Callable[[], TagFunc]
     """
     if not (found := Re.RE_TASK_FMT.search(hook)):
-        raise ValueError("Task does not match with task format regex.")
+        raise ValueError(
+            f"Hook {hook!r} does not match with hook format regex."
+        )
 
     # NOTE: Pass the searching hook string to `path`, `func`, and `tag`.
     hook: HookSearch = HookSearch(**found.groupdict())
@@ -415,7 +490,7 @@ class HookStage(BaseStage):
     Data Validate:
         >>> stage = {
         ...     "name": "Task stage execution",
-        ...     "task": "tasks/function-name@tag-name",
+        ...     "uses": "tasks/function-name@tag-name",
         ...     "args": {
         ...         "FOO": "BAR",
         ...     },
@@ -426,12 +501,13 @@ class HookStage(BaseStage):
         description="A pointer that want to load function from registry.",
     )
     args: DictData = Field(
+        default_factory=dict,
         description="An arguments that want to pass to the hook function.",
         alias="with",
     )
 
     @handler_result()
-    def execute(self, params: DictData) -> DictData:
+    def execute(self, params: DictData) -> Result:
         """Execute the Hook function that already in the hook registry.
 
         :param params: A parameter that want to pass before run any statement.
@@ -440,10 +516,7 @@ class HookStage(BaseStage):
         """
         t_func_hook: str = param2template(self.uses, params)
         t_func: TagFunc = extract_hook(t_func_hook)()
-        if not callable(t_func):
-            raise ImportError(
-                f"Hook caller {t_func_hook!r} function does not callable."
-            )
+
         # VALIDATE: check input task caller parameters that exists before
         #   calling.
         args: DictData = param2template(self.args, params)
@@ -454,7 +527,7 @@ class HookStage(BaseStage):
             if ips.parameters[k].default == Parameter.empty
         ):
             raise ValueError(
-                f"Necessary params, ({', '.join(ips.parameters.keys())}), "
+                f"Necessary params, ({', '.join(ips.parameters.keys())}, ), "
                 f"does not set to args"
             )
         # NOTE: add '_' prefix if it want to use.
@@ -462,9 +535,8 @@ class HookStage(BaseStage):
             if k.removeprefix("_") in args:
                 args[k] = args.pop(k.removeprefix("_"))
 
-        logging.info(
-            f"({self.run_id}) [STAGE]: Hook-Execute: "
-            f"{t_func.name}@{t_func.tag}"
+        logger.info(
+            f"({self.run_id}) [STAGE]: Hook-Execute: {t_func.name}@{t_func.tag}"
         )
         rs: DictData = t_func(**param2template(args, params))
 
@@ -472,11 +544,10 @@ class HookStage(BaseStage):
         #   Check the result type from hook function, it should be dict.
         if not isinstance(rs, dict):
             raise TypeError(
-                f"Return of hook function: {t_func.name}@{t_func.tag} does "
-                f"not serialize to result model, you should fix it to "
-                f"`dict` type."
+                f"Return type: '{t_func.name}@{t_func.tag}' does not serialize "
+                f"to result model, you change return type to `dict`."
             )
-        return rs
+        return Result(status=0, context=rs)
 
 
 class TriggerStage(BaseStage):
@@ -499,8 +570,8 @@ class TriggerStage(BaseStage):
         description="A parameter that want to pass to pipeline execution.",
     )
 
-    @handler_result("Raise from trigger pipeline")
-    def execute(self, params: DictData) -> DictData:
+    @handler_result("Raise from TriggerStage")
+    def execute(self, params: DictData) -> Result:
         """Trigger pipeline execution.
 
         :param params: A parameter data that want to use in this execution.
@@ -510,9 +581,14 @@ class TriggerStage(BaseStage):
 
         # NOTE: Loading pipeline object from trigger name.
         _trigger: str = param2template(self.trigger, params=params)
-        pipe: Pipeline = Pipeline.from_loader(name=_trigger, externals={})
-        rs: Result = pipe.execute(params=param2template(self.params, params))
-        return rs.context
+
+        # NOTE: Set running pipeline ID from running stage ID to external
+        #   params on Loader object.
+        pipe: Pipeline = Pipeline.from_loader(
+            name=_trigger, externals={"run_id": self.run_id}
+        )
+        logger.info(f"({self.run_id}) [STAGE]: Trigger-Execute: {_trigger!r}")
+        return pipe.execute(params=param2template(self.params, params))
 
 
 # NOTE: Order of parsing stage data