from __future__ import annotations import asyncio import typing import datetime from . import app from . import entities as core_entities class TaskContext: """Task tracking context""" current_action: str """Current action being executed""" log: str """Log""" def __init__(self): self.current_action = 'default' self.log = '' def _log(self, msg: str): self.log += msg + '\n' def set_current_action(self, action: str): self.current_action = action def trace( self, msg: str, action: str = None, ): if action is not None: self.set_current_action(action) self._log(f'{datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")} | {self.current_action} | {msg}') def to_dict(self) -> dict: return {'current_action': self.current_action, 'log': self.log} @staticmethod def new() -> TaskContext: return TaskContext() @staticmethod def placeholder() -> TaskContext: global placeholder_context if placeholder_context is None: placeholder_context = TaskContext() return placeholder_context placeholder_context: TaskContext | None = None class TaskWrapper: """Task wrapper""" _id_index: int = 0 """Task ID index""" id: int """Task ID""" task_type: str = 'system' # Task type: system or user """Task type""" kind: str = 'system_task' # Task type determined by the initiator, usually the same task type """Task type""" name: str = '' """Task unique name""" label: str = '' """Task display name""" task_context: TaskContext """Task context""" task: asyncio.Task """Task""" task_stack: list = None """Task stack""" ap: app.Application """Application instance""" scopes: list[core_entities.LifecycleControlScope] """Task scope""" def __init__( self, ap: app.Application, coro: typing.Coroutine, task_type: str = 'system', kind: str = 'system_task', name: str = '', label: str = '', context: TaskContext = None, scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION], ): self.id = TaskWrapper._id_index TaskWrapper._id_index += 1 self.ap = ap self.task_context = context or TaskContext() self.task = self.ap.event_loop.create_task(coro) self.task_type = task_type self.kind = kind self.name = name self.label = label if label != '' else name self.task.set_name(name) self.scopes = scopes def assume_exception(self): try: exception = self.task.exception() if self.task_stack is None: self.task_stack = self.task.get_stack() return exception except Exception: return None def assume_result(self): try: return self.task.result() except Exception: return None def to_dict(self) -> dict: exception_traceback = None if self.assume_exception() is not None: exception_traceback = 'Traceback (most recent call last):\n' for frame in self.task_stack: exception_traceback += ( f' File "{frame.f_code.co_filename}", line {frame.f_lineno}, in {frame.f_code.co_name}\n' ) exception_traceback += f' {self.assume_exception().__str__()}\n' return { 'id': self.id, 'task_type': self.task_type, 'kind': self.kind, 'name': self.name, 'label': self.label, 'scopes': [scope.value for scope in self.scopes], 'task_context': self.task_context.to_dict(), 'runtime': { 'done': self.task.done(), 'state': self.task._state, 'exception': self.assume_exception().__str__() if self.assume_exception() is not None else None, 'exception_traceback': exception_traceback, 'result': self.assume_result().__str__() if self.assume_result() is not None else None, }, } def cancel(self): self.task.cancel() class AsyncTaskManager: """Save all asynchronous tasks in the app Include system-level and user-level (plugin installation, update, etc. initiated by users directly)""" ap: app.Application tasks: list[TaskWrapper] """All tasks""" def __init__(self, ap: app.Application): self.ap = ap self.tasks = [] def create_task( self, coro: typing.Coroutine, task_type: str = 'system', kind: str = 'system-task', name: str = '', label: str = '', context: TaskContext = None, scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION], ) -> TaskWrapper: wrapper = TaskWrapper(self.ap, coro, task_type, kind, name, label, context, scopes) self.tasks.append(wrapper) return wrapper def create_user_task( self, coro: typing.Coroutine, kind: str = 'user-task', name: str = '', label: str = '', context: TaskContext = None, scopes: list[core_entities.LifecycleControlScope] = [core_entities.LifecycleControlScope.APPLICATION], ) -> TaskWrapper: return self.create_task(coro, 'user', kind, name, label, context, scopes) async def wait_all(self): await asyncio.gather(*[t.task for t in self.tasks], return_exceptions=True) def get_all_tasks(self) -> list[TaskWrapper]: return self.tasks def get_tasks_dict( self, type: str = None, ) -> dict: return { 'tasks': [t.to_dict() for t in self.tasks if type is None or t.task_type == type], 'id_index': TaskWrapper._id_index, } def get_task_by_id(self, id: int) -> TaskWrapper | None: for t in self.tasks: if t.id == id: return t return None def cancel_by_scope(self, scope: core_entities.LifecycleControlScope): for wrapper in self.tasks: if not wrapper.task.done() and scope in wrapper.scopes: wrapper.task.cancel() def cancel_task(self, task_id: int): for wrapper in self.tasks: if wrapper.id == task_id: if not wrapper.task.done(): wrapper.task.cancel() return