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

zrb/runner/web_server.py

@@ -0,0 +1,224 @@
+import asyncio
+import datetime
+import json
+import os
+from concurrent.futures import ThreadPoolExecutor
+from functools import partial
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+from threading import Thread
+from typing import Any
+from urllib.parse import parse_qs, urlparse
+
+from ..config import BANNER, SESSION_LOG_DIR, WEB_HTTP_PORT
+from ..context.any_context import AnyContext
+from ..context.shared_context import SharedContext
+from ..group.any_group import AnyGroup
+from ..session.session import Session
+from ..session_state_logger.default_session_state_logger import (
+    default_session_state_logger,
+)
+from ..task.any_task import AnyTask
+from ..util.group import extract_node_from_args, get_node_path
+from .web_app.group_info_ui.controller import handle_group_info_ui
+from .web_app.home_page.controller import handle_home_page
+from .web_app.task_ui.controller import handle_task_ui
+from .web_util import node_path_to_url, start_event_loop, url_to_args
+
+_DIR = os.path.dirname(__file__)
+_STATIC_DIR = os.path.join(_DIR, "web_app", "static")
+executor = ThreadPoolExecutor()
+
+
+class WebRequestHandler(BaseHTTPRequestHandler):
+    def __init__(
+        self,
+        request: Any,
+        client_address: Any,
+        server: Any,
+        root_group: AnyGroup,
+        event_loop: asyncio.AbstractEventLoop,
+        session_dir: str,
+        coros: list,
+    ):
+        self._root_group = root_group
+        self._event_loop = event_loop
+        self._session_dir = session_dir
+        self._coros = coros
+        super().__init__(request=request, client_address=client_address, server=server)
+
+    def do_GET(self):
+        parsed_url = urlparse(self.path)
+        if parsed_url.path in ["/", "/ui", "/ui/"]:
+            handle_home_page(self, self._root_group)
+        elif parsed_url.path == "/pico.min.css":
+            self.send_css_response(os.path.join(_STATIC_DIR, "pico.min.css"))
+        elif parsed_url.path == "/favicon-32x32.png":
+            self.send_image_response(os.path.join(_STATIC_DIR, "favicon-32x32.png"))
+        elif parsed_url.path.startswith("/ui/"):
+            args = url_to_args(parsed_url.path[3:])
+            node, node_path, residual_args = extract_node_from_args(
+                self._root_group, args
+            )
+            url = f"/ui{node_path_to_url(node_path)}"
+            if isinstance(node, AnyTask):
+                shared_ctx = SharedContext(env=dict(os.environ))
+                session = Session(shared_ctx=shared_ctx, root_group=self._root_group)
+                handle_task_ui(
+                    self, self._root_group, node, session, url, residual_args
+                )
+            elif isinstance(node, AnyGroup):
+                handle_group_info_ui(self, self._root_group, node, url)
+            else:
+                self.send_error(404, "Not Found")
+        elif parsed_url.path.startswith("/api/"):
+            args = url_to_args(parsed_url.path[5:])
+            node, _, residual_args = extract_node_from_args(self._root_group, args)
+            if isinstance(node, AnyTask) and len(residual_args) > 0:
+                if residual_args[0] == "list":
+                    task_path = get_node_path(self._root_group, node)
+                    query_params = parse_qs(parsed_url.query)
+                    self.send_session_list_json_data(task_path, query_params)
+                else:
+                    self.send_session_json_data(residual_args[0])
+            else:
+                self.send_error(404, "Not Found")
+        else:
+            self.send_error(404, "Not Found")
+
+    def do_POST(self):
+        parsed_url = urlparse(self.path)
+        if parsed_url.path.startswith("/api/"):
+            args = url_to_args(parsed_url.path[5:])
+            task, _, residual_args = extract_node_from_args(self._root_group, args)
+            session_name = residual_args[0] if len(residual_args) > 0 else None
+            if not isinstance(task, AnyTask):
+                self.send_error(404, "Not found")
+            elif session_name is None:
+                str_kwargs: dict[str, str] = self.read_json_request()
+                shared_ctx = SharedContext(env=dict(os.environ))
+                session = Session(shared_ctx=shared_ctx, root_group=self._root_group)
+                coro = asyncio.run_coroutine_threadsafe(
+                    task.async_run(session, str_kwargs=str_kwargs),
+                    self._event_loop,
+                )
+                self._coros.append(coro)
+                coro.add_done_callback(lambda coro: self._coros.remove(coro))
+                self.send_json_response({"session_name": session.name})
+        else:
+            self.send_error(404, "Not Found")
+
+    def send_session_list_json_data(
+        self, task_path: list[str], query_params: dict[str, list[Any]]
+    ):
+        print(query_params)
+        max_start_time = datetime.datetime.now()
+        if "to" in query_params and len(query_params["to"]) > 0:
+            max_start_time = datetime.datetime.strptime(
+                query_params["to"][0], "%Y-%m-%d %H:%M:%S"
+            )
+        min_start_time = max_start_time - datetime.timedelta(hours=1)
+        if "from" in query_params and len(query_params["from"]) > 0:
+            min_start_time = datetime.datetime.strptime(
+                query_params["from"][0], "%Y-%m-%d %H:%M:%S"
+            )
+        page = 0
+        if "page" in query_params and len(query_params["page"]) > 0:
+            page = int(query_params["page"][0])
+        limit = 10
+        if "limit" in query_params and len(query_params["limit"]) > 0:
+            limit = int(query_params["limit"][0])
+        try:
+            self.send_json_response(
+                default_session_state_logger.list(
+                    task_path,
+                    min_start_time=min_start_time,
+                    max_start_time=max_start_time,
+                    page=page,
+                    limit=limit,
+                )
+            )
+        except Exception as e:
+            self.send_json_response({"error": f"{e}"}, 500)
+
+    def send_session_json_data(self, session_name: str):
+        try:
+            self.send_json_response(default_session_state_logger.read(session_name))
+        except Exception as e:
+            self.send_json_response({"error": f"{e}"}, 500)
+
+    def send_json_response(self, data: Any, http_status: int = 200):
+        self.send_response(http_status)
+        self.send_header("Content-type", "application/json")
+        self.end_headers()
+        self.wfile.write(f"{json.dumps(data)}".encode())
+
+    def send_html_response(self, html: str, http_status: int = 200):
+        self.send_response(http_status)
+        self.send_header("Content-type", "text/html")
+        self.end_headers()
+        self.wfile.write(f"{html}".encode())
+
+    def send_css_response(self, css_path: str):
+        self.send_response(200)
+        self.send_header("Content-type", "text/css")
+        self.end_headers()
+        # If css_content is provided, send it; otherwise, you could read from a file here
+        with open(css_path, "r") as f:
+            css_content = f.read()
+        self.wfile.write(css_content.encode())
+
+    def send_image_response(self, image_path: str):
+        try:
+            with open(image_path, "rb") as img:
+                self.send_response(200)
+                if image_path.endswith(".png"):
+                    self.send_header("Content-type", "image/png")
+                elif image_path.endswith(".jpg") or image_path.endswith(".jpeg"):
+                    self.send_header("Content-type", "image/jpeg")
+                self.end_headers()
+                self.wfile.write(img.read())
+        except FileNotFoundError:
+            self.send_error(404, "Image Not Found")
+
+    def read_json_request(self) -> Any:
+        content_length = int(self.headers["Content-Length"])
+        post_data = self.rfile.read(content_length)
+        try:
+            return json.loads(post_data.decode())
+        except json.JSONDecodeError:
+            self.send_json_response({"error": "Invalid JSON"}, http_status=400)
+            return None
+
+
+def run_web_server(ctx: AnyContext, root_group: AnyGroup, port: int = WEB_HTTP_PORT):
+    server_address = ("", port)
+    event_loop = asyncio.new_event_loop()
+    coros = []
+    # Use functools.partial to bind the custom attribute
+    handler = partial(
+        WebRequestHandler,
+        root_group=root_group,
+        event_loop=event_loop,
+        session_dir=SESSION_LOG_DIR,
+        coros=coros,
+    )
+    httpd = ThreadingHTTPServer(server_address, handler)
+    banner_lines = BANNER.split("\n") + [
+        f"Zrb Server running on http://localhost:{port}"
+    ]
+    for line in banner_lines:
+        ctx.print(line)
+    loop_thread = Thread(target=start_event_loop, args=[event_loop], daemon=True)
+    loop_thread.start()
+    try:
+        httpd.serve_forever()
+    finally:
+        for coro in coros:
+            coro.cancel()
+        httpd.shutdown()
+        httpd.server_close()
+        # Cancel all tasks
+        for task in asyncio.all_tasks(event_loop):
+            task.cancel()
+        event_loop.call_soon_threadsafe(event_loop.stop)
+        loop_thread.join()

zrb/runner/web_util.py

@@ -0,0 +1,46 @@
+import asyncio
+
+
+def url_to_args(url: str) -> list[str]:
+    stripped_url = url.strip("/")
+    return [part for part in stripped_url.split("/") if part.strip() != ""]
+
+
+def node_path_to_url(args: list[str]) -> str:
+    pruned_args = [part for part in args if part.strip() != ""]
+    stripped_url = "/".join(pruned_args)
+    return f"/{stripped_url}/"
+
+
+class SessionSnapshotCondition:
+    def __init__(self):
+        self._should_stop = False
+
+    @property
+    def should_stop(self) -> bool:
+        return self._should_stop
+
+    def stop(self):
+        self._should_stop = True
+
+
+def start_event_loop(event_loop: asyncio.AbstractEventLoop):
+    asyncio.set_event_loop(event_loop)
+    event_loop.run_forever()
+
+
+async def shutdown_event_loop(event_loop: asyncio.AbstractEventLoop):
+    tasks = [
+        t
+        for t in asyncio.all_tasks(event_loop)
+        if not t.done() and t is not asyncio.current_task(event_loop)
+    ]
+    print("TASKS", tasks)
+    for task in tasks:
+        print("CANCEL TASK", task)
+        task.cancel()
+        print("CANCELED", task)
+    await asyncio.gather(*tasks, return_exceptions=True)
+    print("DONE")
+    event_loop.stop()
+    print("LOOP STOP")

zrb/session/any_session.py

@@ -0,0 +1,190 @@
+from __future__ import annotations  # Enables forward references
+
+from abc import ABC, abstractmethod
+from typing import TYPE_CHECKING, Any, Coroutine, TypeVar
+
+from ..context.any_context import AnyContext
+from ..group.any_group import AnyGroup
+from ..session_state_log.session_state_log import SessionStateLog
+from ..session_state_logger.any_session_state_logger import AnySessionStateLogger
+from ..task_status.task_status import TaskStatus
+
+if TYPE_CHECKING:
+    from ..context import any_shared_context
+    from ..task import any_task
+
+TAnySession = TypeVar("TAnySession", bound="AnySession")
+
+
+class AnySession(ABC):
+    """Abstract base class for managing task execution and context in a session.
+
+    This class handles task lifecycle management, context retrieval,
+    deferred task execution, and data exchange between tasks using
+    XCom-like functionality.
+    """
+
+    @property
+    @abstractmethod
+    def name(self) -> str:
+        """Name this session"""
+        pass
+
+    @property
+    @abstractmethod
+    def root_group(self) -> AnyGroup | None:
+        """Session root group"""
+        pass
+
+    @property
+    @abstractmethod
+    def shared_ctx(self) -> any_shared_context.AnySharedContext:
+        """Shared context for this session"""
+        pass
+
+    @abstractmethod
+    def terminate(self):
+        """Terminating session"""
+        pass
+
+    @property
+    @abstractmethod
+    def is_terminated(self) -> bool:
+        """Whether session is terminated or not"""
+        pass
+
+    @property
+    @abstractmethod
+    def parent(self) -> TAnySession | None:
+        """Parent session"""
+        pass
+
+    @abstractmethod
+    def task_path(self) -> str:
+        """Main task's path"""
+        pass
+
+    @property
+    @abstractmethod
+    def final_result(self) -> Any:
+        """Main task's result"""
+        pass
+
+    @property
+    @abstractmethod
+    def state_logger(self) -> AnySessionStateLogger:
+        """State logger"""
+        pass
+
+    @abstractmethod
+    def set_main_task(self, main_task: any_task.AnyTask):
+        """Set main task"""
+        pass
+
+    @abstractmethod
+    def as_state_log(self) -> SessionStateLog:
+        pass
+
+    @abstractmethod
+    def get_ctx(self, task: any_task.AnyTask) -> AnyContext:
+        """Retrieves the context for a specific task.
+
+        Args:
+            task (AnyTask): The task for which to retrieve the context.
+
+        Returns:
+            AnyContext: The context object specific to the provided task.
+        """
+        pass
+
+    @abstractmethod
+    def defer_monitoring(self, task: any_task.AnyTask, coro: Coroutine):
+        """Defers the execution of a task's monitoring coroutine for later processing.
+
+        Args:
+            task (AnyTask): The task associated with the coroutine.
+            coro (Coroutine): The monitoring coroutine to defer.
+        """
+        pass
+
+    @abstractmethod
+    def defer_action(self, task: any_task.AnyTask, coro: Coroutine):
+        """Defers the execution of a task's coroutine for later processing.
+
+        Args:
+            task (AnyTask): The task associated with the coroutine.
+            coro (Coroutine): The coroutine to defer.
+        """
+        pass
+
+    @abstractmethod
+    def defer_coro(self, coro: Coroutine):
+        """Defers the execution of a coroutine for later processing.
+
+        Args:
+            coro (Coroutine): The coroutine to defer.
+        """
+        pass
+
+    @abstractmethod
+    async def wait_deferred(self):
+        """Asynchronously waits for all deffered coroutines to complete"""
+        pass
+
+    @abstractmethod
+    def register_task(self, task: any_task.AnyTask):
+        """Registers a new task in the session.
+
+        Args:
+            task (AnyTask): The task to register in the session.
+        """
+        pass
+
+    @abstractmethod
+    def get_root_tasks(self, task: any_task.AnyTask) -> list[any_task.AnyTask]:
+        """Retrieves the list of root tasks that should be executed first
+        to run the given task.
+
+        Args:
+            task (AnyTask): The current task.
+
+        Returns:
+            list[AnyTask]: A list of root tasks.
+        """
+        pass
+
+    @abstractmethod
+    def get_next_tasks(self, task: any_task.AnyTask) -> list[any_task.AnyTask]:
+        """Retrieves the list of tasks that should be executed after the given task.
+
+        Args:
+            task (AnyTask): The current task.
+
+        Returns:
+            list[AnyTask]: A list of tasks that should be executed next.
+        """
+        pass
+
+    @abstractmethod
+    def get_task_status(self, task: any_task.AnyTask) -> TaskStatus:
+        """Get tasks' status.
+
+        Args:
+            task (AnyTask): The task to mark as started.
+
+        Returns:
+            TaskStatus: Task status object
+        """
+        pass
+
+    @abstractmethod
+    def is_allowed_to_run(self, task: any_task.AnyTask):
+        """Determines if the specified task is allowed to run based on its current state.
+
+        Args:
+            task (AnyTask): The task to check.
+
+        Returns:
+            bool: True if the task is allowed to run, False otherwise.
+        """
+        pass