zrb 1.13.1__py3-none-any.whl → 1.21.17__py3-none-any.whl

zrb/runner/common_util.py

@@ -1,31 +1,36 @@
-from typing import Any
-
 from zrb.context.shared_context import SharedContext
 from zrb.task.any_task import AnyTask
 
 
-def get_run_kwargs(
-    task: AnyTask, args: list[str], kwargs: dict[str, str], cli_mode: bool
+def get_task_str_kwargs(
+    task: AnyTask, str_args: list[str], str_kwargs: dict[str, str], cli_mode: bool
 ) -> dict[str, str]:
     arg_index = 0
-    str_kwargs = {key: f"{val}" for key, val in kwargs.items()}
-    run_kwargs = {**str_kwargs}
-    shared_ctx = SharedContext(args=args)
+    dummmy_shared_ctx = SharedContext()
+    task_str_kwargs = {}
     for task_input in task.inputs:
+        task_name = task_input.name
         if task_input.name in str_kwargs:
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, str_kwargs[task_input.name])
-        elif arg_index < len(args) and task_input.allow_positional_parsing:
-            run_kwargs[task_input.name] = args[arg_index]
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, run_kwargs[task_input.name])
+            task_str_kwargs[task_input.name] = str_kwargs[task_name]
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=str_kwargs[task_name]
+            )
+        elif arg_index < len(str_args) and task_input.allow_positional_parsing:
+            task_str_kwargs[task_name] = str_args[arg_index]
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=task_str_kwargs[task_name]
+            )
             arg_index += 1
         else:
             if cli_mode and task_input.always_prompt:
-                str_value = task_input.prompt_cli_str(shared_ctx)
+                str_value = task_input.prompt_cli_str(dummmy_shared_ctx)
             else:
-                str_value = task_input.get_default_str(shared_ctx)
-            run_kwargs[task_input.name] = str_value
-            # Update shared context for next input default value
-            task_input.update_shared_context(shared_ctx, run_kwargs[task_input.name])
-    return run_kwargs
+                str_value = task_input.get_default_str(dummmy_shared_ctx)
+            task_str_kwargs[task_name] = str_value
+            # Update dummy shared context for next input default value
+            task_input.update_shared_context(
+                dummmy_shared_ctx, str_value=task_str_kwargs[task_name]
+            )
+    return task_str_kwargs

zrb/runner/web_route/task_input_api_route.py

@@ -3,7 +3,7 @@ from typing import TYPE_CHECKING
 
 from zrb.config.web_auth_config import WebAuthConfig
 from zrb.group.any_group import AnyGroup
-from zrb.runner.common_util import get_run_kwargs
+from zrb.runner.common_util import get_task_str_kwargs
 from zrb.runner.web_util.user import get_user_from_request
 from zrb.task.any_task import AnyTask
 from zrb.util.group import NodeNotFoundError, extract_node_from_args
@@ -39,9 +39,9 @@ def serve_task_input_api(
         if isinstance(task, AnyTask):
             if not user.can_access_task(task):
                 return JSONResponse(content={"detail": "Forbidden"}, status_code=403)
-            query_dict = json.loads(query)
-            run_kwargs = get_run_kwargs(
-                task=task, args=[], kwargs=query_dict, cli_mode=False
+            str_kwargs = json.loads(query)
+            task_str_kwargs = get_task_str_kwargs(
+                task=task, str_args=[], str_kwargs=str_kwargs, cli_mode=False
             )
-            return run_kwargs
+            return task_str_kwargs
         return JSONResponse(content={"detail": "Not found"}, status_code=404)

zrb/runner/web_route/task_session_api_route.py

@@ -1,5 +1,4 @@
 import asyncio
-import os
 from datetime import datetime, timedelta
 from typing import TYPE_CHECKING, Any
 
@@ -57,9 +56,7 @@ def serve_task_session_api(
                 return JSONResponse(content={"detail": "Forbidden"}, status_code=403)
             session_name = residual_args[0] if residual_args else None
             if not session_name:
-                shared_ctx = SharedContext(
-                    env={**dict(os.environ), "_ZRB_IS_WEB_MODE": "1"}
-                )
+                shared_ctx = SharedContext(is_web_mode=True)
                 session = Session(shared_ctx=shared_ctx, root_group=root_group)
                 coro = asyncio.create_task(task.async_run(session, str_kwargs=inputs))
                 coroutines.append(coro)

zrb/runner/web_util/user.py

@@ -19,7 +19,7 @@ def get_user_by_credentials(
 
 async def get_user_from_request(
     web_auth_config: WebAuthConfig, request: "Request"
-) -> User | None:
+) -> User:
     from fastapi.security import OAuth2PasswordBearer
 
     if not web_auth_config.enable_auth:
@@ -45,7 +45,11 @@ def _get_user_from_cookie(
     return None
 
 
-def _get_user_from_token(web_auth_config: WebAuthConfig, token: str) -> User | None:
+def _get_user_from_token(
+    web_auth_config: WebAuthConfig, token: str | None
+) -> User | None:
+    if token is None:
+        return None
     try:
         from jose import jwt
 
@@ -54,7 +58,7 @@ def _get_user_from_token(web_auth_config: WebAuthConfig, token: str) -> User | N
             web_auth_config.secret_key,
             options={"require_sub": True, "require_exp": True},
         )
-        username: str = payload.get("sub")
+        username: str | None = payload.get("sub")
         if username is None:
             return None
         user = web_auth_config.find_user_by_username(username)

zrb/session/any_session.py

@@ -1,5 +1,6 @@
 from __future__ import annotations  # Enables forward references
 
+import asyncio
 from abc import ABC, abstractmethod
 from typing import TYPE_CHECKING, Any, Coroutine, TypeVar
 
@@ -62,12 +63,13 @@ class AnySession(ABC):
 
     @property
     @abstractmethod
-    def parent(self) -> TAnySession | None:
+    def parent(self) -> "AnySession | None":
         """Parent session"""
         pass
 
+    @property
     @abstractmethod
-    def task_path(self) -> str:
+    def task_path(self) -> list[str]:
         """Main task's path"""
         pass
 
@@ -105,7 +107,9 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_monitoring(self, task: "AnyTask", coro: Coroutine):
+    def defer_monitoring(
+        self, task: "AnyTask", coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         """Defers the execution of a task's monitoring coroutine for later processing.
 
         Args:
@@ -115,7 +119,9 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_action(self, task: "AnyTask", coro: Coroutine):
+    def defer_action(
+        self, task: "AnyTask", coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         """Defers the execution of a task's coroutine for later processing.
 
         Args:
@@ -125,7 +131,7 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def defer_coro(self, coro: Coroutine):
+    def defer_coro(self, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]):
         """Defers the execution of a coroutine for later processing.
 
         Args:
@@ -185,7 +191,7 @@ class AnySession(ABC):
         pass
 
     @abstractmethod
-    def is_allowed_to_run(self, task: "AnyTask"):
+    def is_allowed_to_run(self, task: "AnyTask") -> bool:
         """Determines if the specified task is allowed to run based on its current state.
 
         Args:

zrb/session/session.py

@@ -1,10 +1,12 @@
+from __future__ import annotations
+
 import asyncio
 from typing import TYPE_CHECKING, Any, Coroutine
 
 from zrb.context.any_shared_context import AnySharedContext
 from zrb.context.context import AnyContext, Context
 from zrb.group.any_group import AnyGroup
-from zrb.session.any_session import AnySession
+from zrb.session.any_session import AnySession, TAnySession
 from zrb.session_state_logger.any_session_state_logger import AnySessionStateLogger
 from zrb.session_state_logger.session_state_logger_factory import session_state_logger
 from zrb.task.any_task import AnyTask
@@ -48,10 +50,10 @@ class Session(AnySession):
         self._context: dict[AnyTask, Context] = {}
         self._shared_ctx = shared_ctx
         self._shared_ctx.set_session(self)
-        self._parent = parent
-        self._action_coros: dict[AnyTask, asyncio.Task] = {}
-        self._monitoring_coros: dict[AnyTask, asyncio.Task] = {}
-        self._coros: list[asyncio.Task] = []
+        self._parent: AnySession | None = parent
+        self._action_coros: dict[AnyTask, asyncio.Task[Any]] = {}
+        self._monitoring_coros: dict[AnyTask, asyncio.Task[Any]] = {}
+        self._coros: list[asyncio.Task[Any]] = []
         self._colors = [
             GREEN,
             YELLOW,
@@ -114,11 +116,13 @@ class Session(AnySession):
         return self._parent
 
     @property
-    def task_path(self) -> str:
+    def task_path(self) -> list[str]:
         return self._main_task_path
 
     @property
     def final_result(self) -> Any:
+        if self._main_task is None:
+            return None
         xcom: Xcom = self.shared_ctx.xcom[self._main_task.name]
         try:
             return xcom.peek()
@@ -134,7 +138,11 @@ class Session(AnySession):
     def set_main_task(self, main_task: AnyTask):
         self.register_task(main_task)
         self._main_task = main_task
-        main_task_path = get_node_path(self._root_group, main_task)
+        main_task_path = (
+            None
+            if self._root_group is None
+            else get_node_path(self._root_group, main_task)
+        )
         self._main_task_path = [] if main_task_path is None else main_task_path
 
     def as_state_log(self) -> "SessionStateLog":
@@ -171,7 +179,7 @@ class Session(AnySession):
         return SessionStateLog(
             name=self.name,
             start_time=log_start_time,
-            main_task_name=self._main_task.name,
+            main_task_name="" if self._main_task is None else self._main_task.name,
             path=self.task_path,
             final_result=(
                 remove_style(f"{self.final_result}")
@@ -188,16 +196,29 @@ class Session(AnySession):
         self._register_single_task(task)
         return self._context[task]
 
-    def defer_monitoring(self, task: AnyTask, coro: Coroutine):
+    def defer_monitoring(
+        self, task: AnyTask, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         self._register_single_task(task)
-        self._monitoring_coros[task] = coro
+        if isinstance(coro, asyncio.Task):
+            self._monitoring_coros[task] = coro
+        else:
+            self._monitoring_coros[task] = asyncio.create_task(coro)
 
-    def defer_action(self, task: AnyTask, coro: Coroutine):
+    def defer_action(
+        self, task: AnyTask, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]
+    ):
         self._register_single_task(task)
-        self._action_coros[task] = coro
+        if isinstance(coro, asyncio.Task):
+            self._action_coros[task] = coro
+        else:
+            self._action_coros[task] = asyncio.create_task(coro)
 
-    def defer_coro(self, coro: Coroutine):
-        self._coros.append(coro)
+    def defer_coro(self, coro: Coroutine[Any, Any, Any] | asyncio.Task[Any]):
+        if isinstance(coro, asyncio.Task):
+            self._coros.append(coro)
+        else:
+            self._coros.append(asyncio.create_task(coro))
         self._coros = [
             existing_coro for existing_coro in self._coros if not existing_coro.done()
         ]
@@ -246,15 +267,15 @@ class Session(AnySession):
 
     def get_next_tasks(self, task: AnyTask) -> list[AnyTask]:
         self._register_single_task(task)
-        return self._downstreams.get(task)
+        return self._downstreams.get(task, [])
 
     def get_task_status(self, task: AnyTask) -> TaskStatus:
         self._register_single_task(task)
         return self._task_status[task]
 
     def _register_single_task(self, task: AnyTask):
-        if task.name not in self._shared_ctx._xcom:
-            self._shared_ctx._xcom[task.name] = Xcom([])
+        if task.name not in self._shared_ctx.xcom:
+            self._shared_ctx.xcom[task.name] = Xcom([])
         if task not in self._context:
             self._context[task] = Context(
                 shared_ctx=self._shared_ctx,
@@ -278,7 +299,7 @@ class Session(AnySession):
             self._color_index = 0
         return chosen
 
-    def _get_icon(self, task: AnyTask) -> int:
+    def _get_icon(self, task: AnyTask) -> str:
         if task.icon is not None:
             return task.icon
         chosen = self._icons[self._icon_index]

zrb/task/any_task.py

@@ -1,7 +1,7 @@
 from __future__ import annotations  # Enables forward references
 
 from abc import ABC, abstractmethod
-from typing import TYPE_CHECKING, Any
+from typing import TYPE_CHECKING, Any, Callable
 
 from zrb.env.any_env import AnyEnv
 from zrb.input.any_input import AnyInput
@@ -36,6 +36,14 @@ class AnyTask(ABC):
     the actual implementation for these abstract members.
     """
 
+    @abstractmethod
+    def __rshift__(self, other: "AnyTask | list[AnyTask]") -> "AnyTask | list[AnyTask]":
+        pass
+
+    @abstractmethod
+    def __lshift__(self, other: "AnyTask | list[AnyTask]") -> "AnyTask":
+        pass
+
     @property
     @abstractmethod
     def name(self) -> str:
@@ -148,13 +156,17 @@ class AnyTask(ABC):
 
     @abstractmethod
     def run(
-        self, session: "AnySession | None" = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: "AnySession | None" = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """Runs the task synchronously.
 
         Args:
             session (AnySession): The shared session.
             str_kwargs(dict[str, str]): The input string values.
+            kwargs(dict[str, Any]): The input values.
 
         Returns:
             Any: The result of the task execution.
@@ -163,13 +175,17 @@ class AnyTask(ABC):
 
     @abstractmethod
     async def async_run(
-        self, session: "AnySession | None" = None, str_kwargs: dict[str, str] = {}
+        self,
+        session: "AnySession | None" = None,
+        str_kwargs: dict[str, str] | None = None,
+        kwargs: dict[str, Any] | None = None,
     ) -> Any:
         """Runs the task asynchronously.
 
         Args:
             session (AnySession): The shared session.
             str_kwargs(dict[str, str]): The input string values.
+            kwargs(dict[str, Any]): The input values.
 
         Returns:
             Any: The result of the task execution.
@@ -203,3 +219,8 @@ class AnyTask(ABC):
             session (AnySession): The shared session.
         """
         pass
+
+    @abstractmethod
+    def to_function(self) -> Callable[..., Any]:
+        """Turn a task into a function"""
+        pass

zrb/task/base/context.py

@@ -1,5 +1,5 @@
 import os
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 from zrb.context.any_context import AnyContext
 from zrb.context.any_shared_context import AnySharedContext
@@ -26,25 +26,33 @@ def build_task_context(task: AnyTask, session: AnySession) -> AnyContext:
 
 
 def fill_shared_context_inputs(
-    task: AnyTask, shared_context: AnySharedContext, str_kwargs: dict[str, str] = {}
+    shared_ctx: AnySharedContext,
+    task: AnyTask,
+    str_kwargs: dict[str, str] | None = None,
+    kwargs: dict[str, Any] | None = None,
 ):
     """
-    Populates the shared context with input values provided via kwargs.
+    Populates the shared context with input values provided via str_kwargs.
     """
+    str_kwarg_dict = str_kwargs if str_kwargs is not None else {}
+    kwarg_dict = kwargs if kwargs is not None else {}
     for task_input in task.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)
+        if task_input.name not in shared_ctx.input:
+            task_input.update_shared_context(
+                shared_ctx,
+                value=kwarg_dict.get(task_input.name, None),
+                str_value=str_kwarg_dict.get(task_input.name, None),
+            )
 
 
-def fill_shared_context_envs(shared_context: AnySharedContext):
+def fill_shared_context_envs(shared_ctx: AnySharedContext):
     """
     Injects OS environment variables into the shared context if they don't already exist.
     """
     os_env_map = {
-        key: val for key, val in os.environ.items() if key not in shared_context.env
+        key: val for key, val in os.environ.items() if key not in shared_ctx.env
     }
-    shared_context.env.update(os_env_map)
+    shared_ctx.env.update(os_env_map)
 
 
 def combine_inputs(

zrb/task/base/execution.py

@@ -53,7 +53,9 @@ def check_execute_condition(task: "BaseTask", session: AnySession) -> bool:
     Evaluates the task's execute_condition attribute.
     """
     ctx = task.get_ctx(session)
-    execute_condition_attr = getattr(task, "_execute_condition", True)
+    execute_condition_attr = (
+        task._execute_condition if task._execute_condition is not None else True
+    )
     return get_bool_attr(ctx, execute_condition_attr, True, auto_render=True)
 
 
@@ -63,8 +65,12 @@ async def execute_action_until_ready(task: "BaseTask", session: AnySession):
     """
     ctx = task.get_ctx(session)
     readiness_checks = task.readiness_checks
-    readiness_check_delay = getattr(task, "_readiness_check_delay", 0.5)
-    monitor_readiness = getattr(task, "_monitor_readiness", False)
+    readiness_check_delay = (
+        task._readiness_check_delay if task._readiness_check_delay is not None else 0.5
+    )
+    monitor_readiness = (
+        task._monitor_readiness if task._monitor_readiness is not None else False
+    )
 
     if not readiness_checks:  # Simplified check for empty list
         ctx.log_info("No readiness checks")
@@ -140,8 +146,8 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
     handling success (triggering successors) and failure (triggering fallbacks).
     """
     ctx = task.get_ctx(session)
-    retries = getattr(task, "_retries", 2)
-    retry_period = getattr(task, "_retry_period", 0)
+    retries = task._retries if task._retries is not None else 2
+    retry_period = task._retry_period if task._retry_period is not None else 0
     max_attempt = retries + 1
     ctx.set_max_attempt(max_attempt)
 
@@ -163,8 +169,9 @@ async def execute_action_with_retry(task: "BaseTask", session: AnySession) -> An
             session.get_task_status(task).mark_as_completed()
 
             # Store result in XCom
-            task_xcom: Xcom = ctx.xcom.get(task.name)
-            task_xcom.push(result)
+            task_xcom: Xcom | None = ctx.xcom.get(task.name)
+            if task_xcom is not None:
+                task_xcom.push(result)
 
             # Skip fallbacks and execute successors on success
             skip_fallbacks(task, session)
@@ -201,7 +208,7 @@ async def run_default_action(task: "BaseTask", ctx: AnyContext) -> Any:
     This is the default implementation called by BaseTask._exec_action.
     Subclasses like LLMTask override _exec_action with their own logic.
     """
-    action = getattr(task, "_action", None)
+    action = task._action
     if action is None:
         ctx.log_debug("No action defined for this task.")
         return None

zrb/task/base/lifecycle.py

@@ -12,7 +12,8 @@ from zrb.util.run import run_async
 async def run_and_cleanup(
     task: AnyTask,
     session: AnySession | None = None,
-    str_kwargs: dict[str, str] = {},
+    str_kwargs: dict[str, str] | None = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> Any:
     """
     Wrapper for async_run that ensures session termination and cleanup of
@@ -23,7 +24,9 @@ async def run_and_cleanup(
         session = Session(shared_ctx=SharedContext())
 
     # Create the main task execution coroutine
-    main_task_coro = asyncio.create_task(run_task_async(task, session, str_kwargs))
+    main_task_coro = asyncio.create_task(
+        run_task_async(task, session, str_kwargs, kwargs)
+    )
 
     try:
         result = await main_task_coro
@@ -67,7 +70,8 @@ async def run_and_cleanup(
 async def run_task_async(
     task: AnyTask,
     session: AnySession | None = None,
-    str_kwargs: dict[str, str] = {},
+    str_kwargs: dict[str, str] | None = None,
+    kwargs: dict[str, Any] | None = None,
 ) -> Any:
     """
     Asynchronous entry point for running a task (`task.async_run()`).
@@ -77,7 +81,7 @@ async def run_task_async(
         session = Session(shared_ctx=SharedContext())
 
     # Populate shared context with inputs and environment variables
-    fill_shared_context_inputs(task, session.shared_ctx, str_kwargs)
+    fill_shared_context_inputs(session.shared_ctx, task, str_kwargs, kwargs)
     fill_shared_context_envs(session.shared_ctx)  # Inject OS env vars
 
     # Start the execution chain from the root tasks

zrb/task/base/monitoring.py

@@ -17,9 +17,13 @@ async def monitor_task_readiness(
     """
     ctx = task.get_ctx(session)
     readiness_checks = task.readiness_checks
-    readiness_check_period = getattr(task, "_readiness_check_period", 5.0)
-    readiness_failure_threshold = getattr(task, "_readiness_failure_threshold", 1)
-    readiness_timeout = getattr(task, "_readiness_timeout", 60)
+    readiness_check_period = (
+        task._readiness_check_period if task._readiness_check_period else 5.0
+    )
+    readiness_failure_threshold = (
+        task._readiness_failure_threshold if task._readiness_failure_threshold else 1
+    )
+    readiness_timeout = task._readiness_timeout if task._readiness_timeout else 60
 
     if not readiness_checks:
         ctx.log_debug("No readiness checks defined, monitoring is not applicable.")
@@ -41,8 +45,9 @@ async def monitor_task_readiness(
                 session.get_task_status(check).reset_history()
                 session.get_task_status(check).reset()
                 # Clear previous XCom data for the check task if needed
-                check_xcom: Xcom = ctx.xcom.get(check.name)
-                check_xcom.clear()
+                check_xcom: Xcom | None = ctx.xcom.get(check.name)
+                if check_xcom is not None:
+                    check_xcom.clear()
 
             readiness_check_coros = [
                 run_async(check.exec_chain(session)) for check in readiness_checks
@@ -77,7 +82,7 @@ async def monitor_task_readiness(
                 )
                 # Ensure check tasks are marked as failed on timeout
                 for check in readiness_checks:
-                    if not session.get_task_status(check).is_finished:
+                    if not session.get_task_status(check).is_ready:
                         session.get_task_status(check).mark_as_failed()
 
             except (asyncio.CancelledError, KeyboardInterrupt):
@@ -92,7 +97,7 @@ async def monitor_task_readiness(
                 )
                 # Mark checks as failed
                 for check in readiness_checks:
-                    if not session.get_task_status(check).is_finished:
+                    if not session.get_task_status(check).is_ready:
                         session.get_task_status(check).mark_as_failed()
 
         # If failure threshold is reached