zrb 0.28.1__py3-none-any.whl → 1.0.0a1__py3-none-any.whl

zrb/session/session.py

@@ -0,0 +1,248 @@
+import asyncio
+from typing import Any, Coroutine
+
+from ..context.any_shared_context import AnySharedContext
+from ..context.context import AnyContext, Context
+from ..group.any_group import AnyGroup
+from ..session_state_log.session_state_log import SessionStateLog, TaskStatusStateLog
+from ..session_state_logger.any_session_state_logger import AnySessionStateLogger
+from ..session_state_logger.default_session_state_logger import (
+    default_session_state_logger,
+)
+from ..task.any_task import AnyTask
+from ..task_status.task_status import TaskStatus
+from ..util.cli.style import BLUE, CYAN, GREEN, ICONS, MAGENTA, YELLOW
+from ..util.group import get_node_path
+from ..util.string.name import get_random_name
+from ..xcom.xcom import Xcom
+from .any_session import AnySession
+
+
+class Session(AnySession):
+    def __init__(
+        self,
+        shared_ctx: AnySharedContext,
+        parent: AnySession | None = None,
+        root_group: AnyGroup | None = None,
+        state_logger: AnySessionStateLogger | None = None,
+    ):
+        self._name = get_random_name()
+        self._root_group = root_group
+        self._state_logger = state_logger
+        self._task_status: dict[AnyTask, TaskStatus] = {}
+        self._upstreams: dict[AnyTask, list[AnyTask]] = {}
+        self._downstreams: dict[AnyTask, list[AnyTask]] = {}
+        self._context: dict[AnyTask, Context] = {}
+        self._shared_ctx = shared_ctx
+        self._shared_ctx.set_session(self)
+        self._parent = parent
+        self._action_coros: dict[AnyTask, Coroutine] = {}
+        self._monitoring_coros: dict[AnyTask, Coroutine] = {}
+        self._coros: list[Coroutine] = []
+        self._colors = [GREEN, YELLOW, BLUE, MAGENTA, CYAN]
+        self._icons = ICONS
+        self._color_index = 0
+        self._icon_index = 0
+        self._is_terminated = False
+        self._main_task: AnyTask | None = None
+        self._main_task_path: list[str] = []
+
+    def __repr__(self):
+        class_name = self.__class__.__name__
+        name = self.name
+        status = self._task_status
+        shared_ctx = self.shared_ctx
+        return f"<{class_name} name={name} status={status}, shared_ctx={shared_ctx}>"
+
+    @property
+    def name(self) -> str:
+        return self._name
+
+    @property
+    def root_group(self) -> AnyGroup | None:
+        return self._root_group
+
+    @property
+    def shared_ctx(self) -> AnySharedContext:
+        return self._shared_ctx
+
+    def terminate(self):
+        self._is_terminated = True
+        for task_status in self._task_status.values():
+            task_status.mark_as_terminated()
+
+    @property
+    def is_terminated(self) -> bool:
+        return self._is_terminated
+
+    @property
+    def parent(self) -> AnySession | None:
+        return self._parent
+
+    @property
+    def task_path(self) -> str:
+        return self._main_task_path
+
+    @property
+    def final_result(self) -> Any:
+        xcom: Xcom = self.shared_ctx.xcom[self._main_task.name]
+        try:
+            return xcom.peek()
+        except IndexError:
+            return None
+
+    @property
+    def state_logger(self) -> AnySessionStateLogger:
+        if self._state_logger is None:
+            return default_session_state_logger
+        return self._state_logger
+
+    def set_main_task(self, main_task: AnyTask):
+        self.register_task(main_task)
+        self._main_task = main_task
+        self._main_task_path = get_node_path(self._root_group, main_task)
+
+    def as_state_log(self) -> SessionStateLog:
+        task_status_log: dict[str, TaskStatusStateLog] = {}
+        for task, task_status in self._task_status.items():
+            task_status_log[task.name] = {
+                "is_started": task_status.is_started,
+                "is_ready": task_status.is_ready,
+                "is_completed": task_status.is_completed,
+                "is_skipped": task_status.is_skipped,
+                "is_failed": task_status.is_failed,
+                "is_permanently_failed": task_status.is_permanently_failed,
+                "is_terminated": task_status.is_terminated,
+                "history": [
+                    {
+                        "status": status,
+                        "time": status_at.strftime("%Y-%m-%d %H:%M:%S.%f"),
+                    }
+                    for status, status_at in task_status.history
+                ],
+            }
+        return {
+            "name": self.name,
+            "main_task_name": self._main_task.name,
+            "path": self.task_path,
+            "final_result": (
+                f"{self.final_result}" if self.final_result is not None else ""
+            ),
+            "finished": self.is_terminated,
+            "log": self.shared_ctx.shared_log,
+            "input": self.shared_ctx.input,
+            "task_status": task_status_log,
+        }
+
+    def get_ctx(self, task: AnyTask) -> AnyContext:
+        self._register_single_task(task)
+        return self._context[task]
+
+    def defer_monitoring(self, task: AnyTask, coro: Coroutine):
+        self._register_single_task(task)
+        self._monitoring_coros[task] = coro
+
+    def defer_action(self, task: AnyTask, coro: Coroutine):
+        self._register_single_task(task)
+        self._action_coros[task] = coro
+
+    def defer_coro(self, coro: Coroutine):
+        self._coros.append(coro)
+        self._coros = [
+            existing_coro for existing_coro in self._coros if not existing_coro.done()
+        ]
+
+    async def wait_deferred(self):
+        await self._wait_deferred_monitoring()
+        await self._wait_deferred_action()
+        await asyncio.gather(*self._coros)
+
+    async def _wait_deferred_action(self):
+        if len(self._action_coros) == 0:
+            return
+        task_coros = self._action_coros.values()
+        await asyncio.gather(*task_coros)
+
+    async def _wait_deferred_monitoring(self):
+        if len(self._monitoring_coros) == 0:
+            return
+        task_coros = self._monitoring_coros.values()
+        await asyncio.gather(*task_coros)
+
+    def register_task(self, task: AnyTask):
+        self._register_single_task(task)
+        for readiness_check in task.readiness_checks:
+            self.register_task(readiness_check)
+        for upstream in task.upstreams:
+            self.register_task(upstream)
+            if task not in self._downstreams[upstream]:
+                self._downstreams[upstream].append(task)
+            if upstream not in self._upstreams[task]:
+                self._upstreams[task].append(upstream)
+
+    def get_root_tasks(self, task: AnyTask) -> list[AnyTask]:
+        root_tasks = []
+        upstreams = self._upstreams[task]
+        if len(upstreams) == 0:
+            root_tasks.append(task)
+        else:
+            for upstream in upstreams:
+                root_tasks += self.get_root_tasks(upstream)
+        return list(set(root_tasks))
+
+    def get_next_tasks(self, task: AnyTask) -> list[AnyTask]:
+        self._register_single_task(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 not in self._context:
+            self._context[task] = Context(
+                shared_ctx=self._shared_ctx,
+                task_name=task.name,
+                color=self._get_color(task),
+                icon=self._get_icon(task),
+            )
+        if task not in self._task_status:
+            self._task_status[task] = TaskStatus()
+        if task not in self._downstreams:
+            self._downstreams[task] = []
+        if task not in self._upstreams:
+            self._upstreams[task] = []
+
+    def _get_color(self, task: AnyTask) -> int:
+        if task.color is not None:
+            return task.color
+        chosen = self._colors[self._color_index]
+        self._color_index += 1
+        if self._color_index >= len(self._colors):
+            self._color_index = 0
+        return chosen
+
+    def _get_icon(self, task: AnyTask) -> int:
+        if task.icon is not None:
+            return task.icon
+        chosen = self._icons[self._icon_index]
+        self._icon_index += 1
+        if self._icon_index >= len(self._icons):
+            self._icon_index = 0
+        return chosen
+
+    def is_allowed_to_run(self, task: AnyTask):
+        if self.is_terminated:
+            return False
+        self._register_single_task(task)
+        task_status = self.get_task_status(task)
+        if task_status.is_started or task_status.is_completed:
+            return False
+        unfulfilled_upstreams = [
+            upstream
+            for upstream in self._upstreams[task]
+            if not self._task_status[upstream].allow_run_downstream
+        ]
+        return len(unfulfilled_upstreams) == 0

zrb/session_state_log/session_state_log.py

@@ -0,0 +1,34 @@
+from typing import Any, TypedDict
+
+
+class TaskStatusHistoryStateLog(TypedDict):
+    status: str
+    time: str
+
+
+class TaskStatusStateLog(TypedDict):
+    history: list[TaskStatusHistoryStateLog]
+    is_started: bool
+    is_ready: bool
+    is_completed: bool
+    is_skipped: bool
+    is_failed: bool
+    is_permanently_failed: bool
+    is_terminated: bool
+
+
+class SessionStateLog(TypedDict):
+    name: str
+    start_time: str
+    main_task_name: str
+    path: list[str]
+    input: dict[str, Any]
+    final_result: str
+    finished: bool
+    log: list[str]
+    task_status: dict[str, TaskStatusStateLog]
+
+
+class SessionStateLogList(TypedDict):
+    total: int
+    data: list[SessionStateLog]

zrb/session_state_logger/any_session_state_logger.py

@@ -0,0 +1,26 @@
+import datetime
+from abc import ABC, abstractmethod
+
+from ..session_state_log.session_state_log import SessionStateLog, SessionStateLogList
+
+
+class AnySessionStateLogger(ABC):
+
+    @abstractmethod
+    def write(self, session_log: SessionStateLog):
+        pass
+
+    @abstractmethod
+    def read(self, session_name: str) -> SessionStateLog:
+        pass
+
+    @abstractmethod
+    def list(
+        self,
+        task_path: list[str],
+        min_start_time: datetime.datetime,
+        max_start_time: datetime.datetime,
+        page: int = 0,
+        limit: int = 10,
+    ) -> SessionStateLogList:
+        pass

zrb/session_state_logger/default_session_state_logger.py

@@ -0,0 +1,4 @@
+from ..config import SESSION_LOG_DIR
+from .file_session_state_logger import FileSessionStateLogger
+
+default_session_state_logger = FileSessionStateLogger(SESSION_LOG_DIR)

zrb/session_state_logger/file_session_state_logger.py

@@ -0,0 +1,98 @@
+import datetime
+import json
+import os
+
+from ..session_state_log.session_state_log import SessionStateLog, SessionStateLogList
+from .any_session_state_logger import AnySessionStateLogger
+
+
+class FileSessionStateLogger(AnySessionStateLogger):
+
+    def __init__(self, session_log_dir: str):
+        self._session_log_dir = session_log_dir
+
+    def write(self, session_log: SessionStateLog):
+        session_file_path = self._get_session_file_path(session_log["name"])
+        with open(session_file_path, "w") as f:
+            f.write(json.dumps(session_log))
+        start_time = self._get_start_time(session_log)
+        if start_time is None:
+            return
+        timeline_dir_path = self._get_timeline_dir_path(session_log, start_time)
+        os.makedirs(timeline_dir_path, exist_ok=True)
+        with open(os.path.join(timeline_dir_path, session_log["name"]), "w"):
+            pass
+
+    def read(self, session_name: str) -> SessionStateLog:
+        session_file_path = self._get_session_file_path(session_name)
+        with open(session_file_path, "r") as f:
+            return json.loads(f.read())
+
+    def list(
+        self,
+        task_path: list[str],
+        min_start_time: datetime.datetime,
+        max_start_time: datetime.datetime,
+        page: int = 0,
+        limit: int = 10,
+    ) -> SessionStateLogList:
+        matching_sessions = []
+        # Traverse the timeline directory and filter sessions
+        timeline_dir = os.path.join(self._session_log_dir, "_timeline", *task_path)
+        if not os.path.exists(timeline_dir):
+            return {"total": 0, "data": []}
+        for root, _, files in os.walk(timeline_dir):
+            for file_name in files:
+                session_name = os.path.splitext(file_name)[0]
+                # Read the session and retrieve start time
+                session_log = self.read(session_name)
+                start_time = self._get_start_time(session_log)
+                # Filter sessions based on start time
+                if start_time and min_start_time <= start_time <= max_start_time:
+                    matching_sessions.append((start_time, session_log))
+        # Sort sessions by start time, descending
+        matching_sessions.sort(key=lambda x: x[0], reverse=True)
+        total = len(matching_sessions)
+        # Apply pagination
+        start_index = page * limit
+        end_index = start_index + limit
+        paginated_sessions = matching_sessions[start_index:end_index]
+        # Extract session logs from the sorted list of tuples
+        data = [session_log for _, session_log in paginated_sessions]
+        return {"total": total, "data": data}
+
+    def _get_session_file_path(self, session_name: str) -> str:
+        return os.path.join(self._session_log_dir, f"{session_name}.json")
+
+    def _get_timeline_dir_path(
+        self, session_log: SessionStateLog, start_time: datetime.datetime
+    ) -> str:
+        year = start_time.year
+        month = start_time.month
+        day = start_time.day
+        hour = start_time.hour
+        minute = start_time.minute
+        second = start_time.second
+        paths = session_log["path"] + [
+            f"{year}",
+            f"{month}",
+            f"{day}",
+            f"{hour}",
+            f"{minute}",
+            f"{second}",
+        ]
+        return os.path.join(self._session_log_dir, "_timeline", *paths)
+
+    def _get_start_time(self, session_log: SessionStateLog) -> datetime.datetime:
+        result: datetime.datetime | None = None
+        for task_status in session_log["task_status"].values():
+            histories = task_status["history"]
+            if len(histories) == 0:
+                continue
+            first_history = histories[0]
+            first_time = datetime.datetime.strptime(
+                first_history["time"], "%Y-%m-%d %H:%M:%S.%f"
+            )
+            if result is None or first_time < result:
+                result = first_time
+        return result