zrb 1.5.7__py3-none-any.whl → 1.5.9__py3-none-any.whl

zrb/task/base_task.py

@@ -1,19 +1,25 @@
 import asyncio
-import os
-from collections.abc import Callable, Coroutine
+from collections.abc import Callable
 from typing import Any
 
 from zrb.attr.type import BoolAttr, fstring
 from zrb.context.any_context import AnyContext
-from zrb.context.shared_context import AnySharedContext, SharedContext
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
 from zrb.session.any_session import AnySession
-from zrb.session.session import Session
 from zrb.task.any_task import AnyTask
-from zrb.util.attr import get_bool_attr
-from zrb.util.run import run_async
-from zrb.xcom.xcom import Xcom
+from zrb.task.base.context import (
+    build_task_context,
+    get_combined_envs,
+    get_combined_inputs,
+)
+from zrb.task.base.execution import (
+    execute_task_action,
+    execute_task_chain,
+    run_default_action,
+)
+from zrb.task.base.lifecycle import execute_root_tasks, run_and_cleanup, run_task_async
+from zrb.task.base.operators import handle_lshift, handle_rshift
 
 
 class BaseTask(AnyTask):
@@ -62,25 +68,13 @@ class BaseTask(AnyTask):
         self._action = action
 
     def __repr__(self):
-        return f"<{self.__class__.__name__} name={self._name}>"
+        return f"<{self.__class__.__name__} name={self.name}>"
 
     def __rshift__(self, other: AnyTask | list[AnyTask]) -> AnyTask | list[AnyTask]:
-        try:
-            if isinstance(other, AnyTask):
-                other.append_upstream(self)
-            elif isinstance(other, list):
-                for task in other:
-                    task.append_upstream(self)
-            return other
-        except Exception as e:
-            raise ValueError(f"Invalid operation {self} >> {other}: {e}")
+        return handle_rshift(self, other)
 
     def __lshift__(self, other: AnyTask | list[AnyTask]) -> AnyTask:
-        try:
-            self.append_upstream(other)
-            return self
-        except Exception as e:
-            raise ValueError(f"Invalid operation {self} << {other}: {e}")
+        return handle_lshift(self, other)
 
     @property
     def name(self) -> str:
@@ -104,430 +98,130 @@ class BaseTask(AnyTask):
 
     @property
     def envs(self) -> list[AnyEnv]:
-        envs = []
-        for upstream in self.upstreams:
-            envs += upstream.envs
-        if isinstance(self._envs, AnyEnv):
-            envs.append(self._envs)
-        elif self._envs is not None:
-            envs += self._envs
-        return [env for env in envs if env is not None]
+        return get_combined_envs(self)
 
     @property
     def inputs(self) -> list[AnyInput]:
-        inputs = []
-        for upstream in self.upstreams:
-            self.__combine_inputs(inputs, upstream.inputs)
-        if self._inputs is not None:
-            self.__combine_inputs(inputs, self._inputs)
-        return [task_input for task_input in inputs if task_input is not None]
-
-    def __combine_inputs(
-        self,
-        inputs: list[AnyInput],
-        other_inputs: list[AnyInput | None] | AnyInput | None,
-    ):
-        input_names = [task_input.name for task_input in inputs]
-        if isinstance(other_inputs, AnyInput):
-            other_inputs = [other_inputs]
-        elif other_inputs is None:
-            other_inputs = []
-        for task_input in other_inputs:
-            if task_input is None:
-                continue
-            if task_input.name not in input_names:
-                inputs.append(task_input)
+        return get_combined_inputs(self)
 
     @property
     def fallbacks(self) -> list[AnyTask]:
+        """Returns the list of fallback tasks."""
         if self._fallbacks is None:
             return []
-        elif isinstance(self._fallbacks, AnyTask):
-            return [self._fallbacks]
-        return self._fallbacks
+        elif isinstance(self._fallbacks, list):
+            return self._fallbacks
+        return [self._fallbacks]  # Assume single task
 
     def append_fallback(self, fallbacks: AnyTask | list[AnyTask]):
-        fallback_list = [fallbacks] if isinstance(fallbacks, AnyTask) else fallbacks
-        for fallback in fallback_list:
-            self.__append_fallback(fallback)
-
-    def __append_fallback(self, fallback: AnyTask):
-        # Make sure self._fallbacks is a list
+        """Appends fallback tasks, ensuring no duplicates."""
         if self._fallbacks is None:
             self._fallbacks = []
-        elif isinstance(self._fallbacks, AnyTask):
+        elif not isinstance(self._fallbacks, list):
             self._fallbacks = [self._fallbacks]
-        # Add fallback if it was not on self._fallbacks
-        if fallback not in self._fallbacks:
-            self._fallbacks.append(fallback)
+        to_add = fallbacks if isinstance(fallbacks, list) else [fallbacks]
+        for fb in to_add:
+            if fb not in self._fallbacks:
+                self._fallbacks.append(fb)
 
     @property
     def successors(self) -> list[AnyTask]:
+        """Returns the list of successor tasks."""
         if self._successors is None:
             return []
-        elif isinstance(self._successors, AnyTask):
-            return [self._successors]
-        return self._successors
+        elif isinstance(self._successors, list):
+            return self._successors
+        return [self._successors]  # Assume single task
 
     def append_successor(self, successors: AnyTask | list[AnyTask]):
-        successor_list = [successors] if isinstance(successors, AnyTask) else successors
-        for successor in successor_list:
-            self.__append_successor(successor)
-
-    def __append_successor(self, successor: AnyTask):
-        # Make sure self._successors is a list
+        """Appends successor tasks, ensuring no duplicates."""
         if self._successors is None:
             self._successors = []
-        elif isinstance(self._successors, AnyTask):
+        elif not isinstance(self._successors, list):
             self._successors = [self._successors]
-        # Add successor if it was not on self._successors
-        if successor not in self._successors:
-            self._successors.append(successor)
+        to_add = successors if isinstance(successors, list) else [successors]
+        for succ in to_add:
+            if succ not in self._successors:
+                self._successors.append(succ)
 
     @property
     def readiness_checks(self) -> list[AnyTask]:
+        """Returns the list of readiness check tasks."""
         if self._readiness_checks is None:
             return []
-        elif isinstance(self._readiness_checks, AnyTask):
-            return [self._readiness_checks]
-        return self._readiness_checks
+        elif isinstance(self._readiness_checks, list):
+            return self._readiness_checks
+        return [self._readiness_checks]  # Assume single task
 
     def append_readiness_check(self, readiness_checks: AnyTask | list[AnyTask]):
-        readiness_check_list = (
-            [readiness_checks]
-            if isinstance(readiness_checks, AnyTask)
-            else readiness_checks
-        )
-        for readiness_check in readiness_check_list:
-            self.__append_readiness_check(readiness_check)
-
-    def __append_readiness_check(self, readiness_check: AnyTask):
-        # Make sure self._readiness_checks is a list
+        """Appends readiness check tasks, ensuring no duplicates."""
         if self._readiness_checks is None:
             self._readiness_checks = []
-        elif isinstance(self._readiness_checks, AnyTask):
+        elif not isinstance(self._readiness_checks, list):
             self._readiness_checks = [self._readiness_checks]
-        # Add readiness_check if it was not on self._readiness_checks
-        if readiness_check not in self._readiness_checks:
-            self._readiness_checks.append(readiness_check)
+        to_add = (
+            readiness_checks
+            if isinstance(readiness_checks, list)
+            else [readiness_checks]
+        )
+        for rc in to_add:
+            if rc not in self._readiness_checks:
+                self._readiness_checks.append(rc)
 
     @property
     def upstreams(self) -> list[AnyTask]:
+        """Returns the list of upstream tasks."""
         if self._upstreams is None:
             return []
-        elif isinstance(self._upstreams, AnyTask):
-            return [self._upstreams]
-        return self._upstreams
+        elif isinstance(self._upstreams, list):
+            return self._upstreams
+        return [self._upstreams]  # Assume single task
 
     def append_upstream(self, upstreams: AnyTask | list[AnyTask]):
-        upstream_list = [upstreams] if isinstance(upstreams, AnyTask) else upstreams
-        for upstream in upstream_list:
-            self.__append_upstream(upstream)
-
-    def __append_upstream(self, upstream: AnyTask):
-        # Make sure self._upstreams is a list
+        """Appends upstream tasks, ensuring no duplicates."""
         if self._upstreams is None:
             self._upstreams = []
-        elif isinstance(self._upstreams, AnyTask):
+        elif not isinstance(self._upstreams, list):
             self._upstreams = [self._upstreams]
-        # Add upstream if it was not on self._upstreams
-        if upstream not in self._upstreams:
-            self._upstreams.append(upstream)
+        to_add = upstreams if isinstance(upstreams, list) else [upstreams]
+        for up in to_add:
+            if up not in self._upstreams:
+                self._upstreams.append(up)
 
     def get_ctx(self, session: AnySession) -> AnyContext:
-        ctx = session.get_ctx(self)
-        # Enhance session ctx with current task env
-        for env in self.envs:
-            env.update_context(ctx)
-        return ctx
+        return build_task_context(self, session)
 
     def run(
         self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
     ) -> Any:
-        return asyncio.run(self._run_and_cleanup(session, str_kwargs))
-
-    async def _run_and_cleanup(
-        self,
-        session: AnySession | None = None,
-        str_kwargs: dict[str, str] = {},
-    ) -> Any:
-        current_task = asyncio.create_task(self.async_run(session, str_kwargs))
-        try:
-            result = await current_task
-        finally:
-            if session and not session.is_terminated:
-                session.terminate()
-            # Cancel all running tasks except the current one
-            pending = [task for task in asyncio.all_tasks() if task is not current_task]
-            for task in pending:
-                task.cancel()
-            if pending:
-                try:
-                    await asyncio.wait(pending, timeout=5)
-                except asyncio.CancelledError:
-                    pass
-        return result
+        # Use asyncio.run() to execute the async cleanup wrapper
+        return asyncio.run(run_and_cleanup(self, session, str_kwargs))
 
     async def async_run(
         self, session: AnySession | None = None, str_kwargs: dict[str, str] = {}
     ) -> Any:
-        if session is None:
-            session = Session(shared_ctx=SharedContext())
-        # Update session
-        self.__fill_shared_context_inputs(session.shared_ctx, str_kwargs)
-        self.__fill_shared_context_envs(session.shared_ctx)
-        result = await run_async(self.exec_root_tasks(session))
-        return result
-
-    def __fill_shared_context_inputs(
-        self, shared_context: AnySharedContext, str_kwargs: dict[str, str] = {}
-    ):
-        for task_input in self.inputs:
-            if task_input.name not in shared_context.input:
-                str_value = str_kwargs.get(task_input.name, None)
-                task_input.update_shared_context(shared_context, str_value)
-
-    def __fill_shared_context_envs(self, shared_context: AnySharedContext):
-        # Inject os environ
-        os_env_map = {
-            key: val for key, val in os.environ.items() if key not in shared_context.env
-        }
-        shared_context.env.update(os_env_map)
+        return await run_task_async(self, session, str_kwargs)
 
     async def exec_root_tasks(self, session: AnySession):
-        session.set_main_task(self)
-        session.state_logger.write(session.as_state_log())
-        try:
-            log_state = asyncio.create_task(self._log_session_state(session))
-            root_tasks = [
-                task
-                for task in session.get_root_tasks(self)
-                if session.is_allowed_to_run(task)
-            ]
-            root_task_coros = [
-                run_async(root_task.exec_chain(session)) for root_task in root_tasks
-            ]
-            await asyncio.gather(*root_task_coros)
-            await session.wait_deferred()
-            session.terminate()
-            await log_state
-            return session.final_result
-        except IndexError:
-            return None
-        except (asyncio.CancelledError, KeyboardInterrupt):
-            ctx = self.get_ctx(session)
-            ctx.log_info("Session terminated")
-        finally:
-            session.terminate()
-            session.state_logger.write(session.as_state_log())
-            ctx = self.get_ctx(session)
-            ctx.log_debug(session)
-
-    async def _log_session_state(self, session: AnySession):
-        session.state_logger.write(session.as_state_log())
-        while not session.is_terminated:
-            await asyncio.sleep(0.1)
-            session.state_logger.write(session.as_state_log())
-        session.state_logger.write(session.as_state_log())
+        return await execute_root_tasks(self, session)
 
     async def exec_chain(self, session: AnySession):
-        if session.is_terminated or not session.is_allowed_to_run(self):
-            return
-        result = await self.exec(session)
-        # Get next tasks
-        nexts = session.get_next_tasks(self)
-        if session.is_terminated or len(nexts) == 0:
-            return result
-        # Run next tasks asynchronously
-        next_coros = [run_async(next.exec_chain(session)) for next in nexts]
-        return await asyncio.gather(*next_coros)
+        return await execute_task_chain(self, session)
 
     async def exec(self, session: AnySession):
-        ctx = self.get_ctx(session)
-        if not session.is_allowed_to_run(self):
-            # Task is not allowed to run, skip it for now.
-            # This will be triggered later
-            ctx.log_info("Not allowed to run")
-            return
-        if not self.__get_execute_condition(session):
-            # Skip the task
-            ctx.log_info("Marked as skipped")
-            session.get_task_status(self).mark_as_skipped()
-            return
-        # Wait for task to be ready
-        return await run_async(self.__exec_action_until_ready(session))
-
-    def __get_execute_condition(self, session: Session) -> bool:
-        ctx = self.get_ctx(session)
-        return get_bool_attr(ctx, self._execute_condition, True, auto_render=True)
-
-    async def __exec_action_until_ready(self, session: AnySession):
-        ctx = self.get_ctx(session)
-        readiness_checks = self.readiness_checks
-        if len(readiness_checks) == 0:
-            ctx.log_info("No readiness checks")
-            # Task has no readiness check
-            result = await run_async(self.__exec_action_and_retry(session))
-            ctx.log_info("Marked as ready")
-            session.get_task_status(self).mark_as_ready()
-            return result
-        # Start the task along with the readiness checks
-        action_coro = asyncio.create_task(
-            run_async(self.__exec_action_and_retry(session))
-        )
-        await asyncio.sleep(self._readiness_check_delay)
-        readiness_check_coros = [
-            run_async(check.exec_chain(session)) for check in readiness_checks
-        ]
-        # Only wait for readiness checks and mark the task as ready
-        ctx.log_info("Start readiness checks")
-        result = await asyncio.gather(*readiness_check_coros)
-        ctx.log_info("Readiness checks completed")
-        ctx.log_info("Marked as ready")
-        session.get_task_status(self).mark_as_ready()
-        # Defer task's coroutines, will be waited later
-        session.defer_action(self, action_coro)
-        if self._monitor_readiness:
-            monitor_and_rerun_coro = asyncio.create_task(
-                run_async(self.__exec_monitoring(session, action_coro))
-            )
-            session.defer_monitoring(self, monitor_and_rerun_coro)
-        return result
-
-    async def __exec_monitoring(self, session: AnySession, action_coro: asyncio.Task):
-        readiness_checks = self.readiness_checks
-        failure_count = 0
-        ctx = self.get_ctx(session)
-        while not session.is_terminated:
-            await asyncio.sleep(self._readiness_check_period)
-            if failure_count < self._readiness_failure_threshold:
-                for readiness_check in readiness_checks:
-                    session.get_task_status(readiness_check).reset_history()
-                    session.get_task_status(readiness_check).reset()
-                    readiness_xcom: Xcom = ctx.xcom[self.name]
-                    readiness_xcom.clear()
-                readiness_check_coros = [
-                    check.exec_chain(session) for check in readiness_checks
-                ]
-                try:
-                    ctx.log_info("Checking")
-                    await asyncio.wait_for(
-                        asyncio.gather(*readiness_check_coros),
-                        timeout=self._readiness_timeout,
-                    )
-                    ctx.log_info("OK")
-                    continue
-                except (asyncio.CancelledError, KeyboardInterrupt):
-                    for readiness_check in readiness_checks:
-                        ctx.log_info("Marked as failed")
-                        session.get_task_status(readiness_check).mark_as_failed()
-                except asyncio.TimeoutError:
-                    failure_count += 1
-                    ctx.log_info("Detecting failure")
-                    ctx.log_debug(f"Failure count: {failure_count}")
-            # Readiness check failed, reset
-            ctx.log_info("Resetting")
-            action_coro.cancel()
-            session.get_task_status(self).reset()
-            # defer this action
-            ctx.log_info("Running")
-            action_coro: Coroutine = asyncio.create_task(
-                run_async(self.__exec_action_and_retry(session))
-            )
-            session.defer_action(self, action_coro)
-            failure_count = 0
-            ctx.log_info("Continue monitoring")
-
-    async def __exec_action_and_retry(self, session: AnySession) -> Any:
-        """
-        Executes an action with retry logic.
-
-        This method attempts to execute the action defined in `_exec_action` with a specified number of retries.
-        If the action fails, it will retry after a specified period until the maximum number of attempts is reached.
-        If the action succeeds, it marks the task as completed and executes any successors.
-        If the action fails permanently, it marks the task as permanently failed and executes any fallbacks.
-
-        Args:
-            session (AnySession): The session object containing the task status and context.
-
-        Returns:
-            Any: The result of the executed action if successful.
-
-        Raises:
-            Exception: If the action fails permanently after all retry attempts.
-        """
-        ctx = self.get_ctx(session)
-        max_attempt = self._retries + 1
-        ctx.set_max_attempt(max_attempt)
-        for attempt in range(max_attempt):
-            ctx.set_attempt(attempt + 1)
-            if attempt > 0:
-                # apply retry period only if this is not the first attempt
-                await asyncio.sleep(self._retry_period)
-            try:
-                ctx.log_info("Marked as started")
-                session.get_task_status(self).mark_as_started()
-                result = await run_async(self._exec_action(ctx))
-                ctx.log_info("Marked as completed")
-                session.get_task_status(self).mark_as_completed()
-                # Put result on xcom
-                task_xcom: Xcom = ctx.xcom.get(self.name)
-                task_xcom.push(result)
-                self.__skip_fallbacks(session)
-                await run_async(self.__exec_successors(session))
-                return result
-            except (asyncio.CancelledError, KeyboardInterrupt):
-                ctx.log_info("Marked as failed")
-                session.get_task_status(self).mark_as_failed()
-                return
-            except BaseException as e:
-                ctx.log_error(e)
-                if attempt < max_attempt - 1:
-                    ctx.log_info("Marked as failed")
-                    session.get_task_status(self).mark_as_failed()
-                    continue
-                ctx.log_info("Marked as permanently failed")
-                session.get_task_status(self).mark_as_permanently_failed()
-                self.__skip_successors(session)
-                await run_async(self.__exec_fallbacks(session))
-                raise e
-
-    async def __exec_successors(self, session: AnySession) -> Any:
-        successors: list[AnyTask] = self.successors
-        successor_coros = [
-            run_async(successor.exec_chain(session)) for successor in successors
-        ]
-        await asyncio.gather(*successor_coros)
-
-    def __skip_successors(self, session: AnySession) -> Any:
-        for successor in self.successors:
-            session.get_task_status(successor).mark_as_skipped()
-
-    async def __exec_fallbacks(self, session: AnySession) -> Any:
-        fallbacks: list[AnyTask] = self.fallbacks
-        fallback_coros = [
-            run_async(fallback.exec_chain(session)) for fallback in fallbacks
-        ]
-        await asyncio.gather(*fallback_coros)
-
-    def __skip_fallbacks(self, session: AnySession) -> Any:
-        for fallback in self.fallbacks:
-            session.get_task_status(fallback).mark_as_skipped()
+        return await execute_task_action(self, session)
 
     async def _exec_action(self, ctx: AnyContext) -> Any:
-        """Execute the main action of the task.
-        By default will render and run the _action attribute.
-
-        Override this method to define custom action.
+        """
+        Execute the main action of the task.
+        This is the primary method to override in subclasses for custom action logic.
+        The default implementation handles the '_action' attribute (string or callable).
 
         Args:
-            session (AnySession): The shared session.
+            ctx (AnyContext): The execution context for this task.
 
         Returns:
             Any: The result of the action execution.
         """
-        if self._action is None:
-            return
-        if isinstance(self._action, str):
-            return ctx.render(self._action)
-        return await run_async(self._action(ctx))
+        # Delegate to the helper function for the default behavior
+        return await run_default_action(self, ctx)

zrb/task/cmd_task.py

@@ -130,7 +130,8 @@ class CmdTask(BaseTask):
             partial(ctx.print, plain=True) if self._should_plain_print else ctx.print
         )
         xcom_pid_key = f"{self.name}-pid"
-        ctx.xcom[xcom_pid_key] = Xcom([])
+        if xcom_pid_key not in ctx.xcom:
+            ctx.xcom[xcom_pid_key] = Xcom([])
         cmd_result, return_code = await run_command(
             cmd=[shell, shell_flag, cmd_script],
             cwd=cwd,