telegrinder 1.0.0rc1__py3-none-any.whl

telegrinder/tools/lifespan.py

@@ -0,0 +1,188 @@
+import asyncio
+import dataclasses
+import datetime
+import typing
+from contextlib import asynccontextmanager
+
+from telegrinder.modules import logger
+from telegrinder.tools.aio import run_task
+from telegrinder.tools.fullname import fullname
+
+type CoroutineTask[T] = typing.Coroutine[typing.Any, typing.Any, T]
+type CoroutineFunc[**P, T] = typing.Callable[P, CoroutineTask[T]]
+type Task[**P, T] = CoroutineFunc[P, T] | CoroutineTask[T] | DelayedTask[P]
+
+
+def to_coroutine_task[T](task: Task[..., T], /) -> CoroutineTask[T]:
+    if asyncio.iscoroutinefunction(task) or isinstance(task, DelayedTask):
+        task = task()
+    elif not asyncio.iscoroutine(task):
+        raise TypeError("Task should be coroutine or coroutine function.")
+    return task
+
+
+@dataclasses.dataclass
+class DelayedTask[**P]:
+    _cancelled: bool = dataclasses.field(default=False, init=False, repr=False)
+    _event: asyncio.Event = dataclasses.field(default_factory=asyncio.Event, init=False, repr=False)
+    _timer: asyncio.TimerHandle | None = dataclasses.field(default=None, init=False, repr=False)
+
+    function: typing.Callable[P, CoroutineTask[typing.Any]]
+    seconds: float | datetime.timedelta
+    repeat: bool = dataclasses.field(default=False, kw_only=True)
+
+    def __post_init__(self) -> None:
+        self.function.cancel = self.cancel
+
+    async def __call__(self, *args: P.args, **kwargs: P.kwargs) -> None:
+        stopped = False
+
+        while not stopped and not self.is_cancelled:
+            self.start_timer()
+            await self._event.wait()
+
+            try:
+                await self.function(*args, **kwargs)
+            except Exception:
+                await logger.aexception(
+                    "Delayed task `{}` failed with exception, traceback message below:",
+                    fullname(self.function),
+                )
+            finally:
+                self._event.clear()
+                self._timer = None
+
+                if not self.repeat:
+                    stopped = True
+
+    @property
+    def is_cancelled(self) -> bool:
+        return self._cancelled
+
+    @property
+    def delay(self) -> float:
+        return float(self.seconds) if isinstance(self.seconds, int | float) else self.seconds.total_seconds()
+
+    def start_timer(self) -> None:
+        if self._timer is None:
+            self._timer = asyncio.get_running_loop().call_later(
+                self.delay,
+                callback=lambda: (self._event.set() if not self._event.is_set() else None),
+            )
+
+    def cancel(self) -> bool:
+        if self._cancelled:
+            return True
+
+        self._cancelled = True
+
+        if self._timer is None or self._timer.cancelled():
+            self._timer = None
+            return False
+
+        self._timer.cancel()
+        self._timer = None
+
+        for future in self._event._waiters:
+            if not future.cancelled():
+                future.cancel()
+
+        return True
+
+
+@dataclasses.dataclass(kw_only=True, slots=True, repr=False)
+class Lifespan:
+    _started: bool = dataclasses.field(default=False, init=False)
+    _lifespan_context: typing.AsyncContextManager[typing.Any] | None = dataclasses.field(default=None, init=False)
+    lifespan_function: typing.Callable[[], typing.AsyncContextManager[typing.Any]] | None = dataclasses.field(
+        default=None
+    )
+    startup_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
+    shutdown_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
+
+    def __repr__(self) -> str:
+        return "<{}: started={}>".format(fullname(self), self._started)
+
+    def __add__(self, other: object, /) -> typing.Self:
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+
+        return self.__class__(
+            startup_tasks=self.startup_tasks + other.startup_tasks,
+            shutdown_tasks=self.shutdown_tasks + other.shutdown_tasks,
+        )
+
+    def __iadd__(self, other: object, /) -> typing.Self:
+        if not isinstance(other, self.__class__):
+            return NotImplemented
+
+        self.startup_tasks.extend(other.startup_tasks)
+        self.shutdown_tasks.extend(other.shutdown_tasks)
+        return self
+
+    async def __aenter__(self) -> None:
+        await self._start()
+
+    async def __aexit__(
+        self,
+        exc_type: typing.Any | None,
+        exc_value: typing.Any | None,
+        exc_tb: typing.Any | None,
+    ) -> None:
+        await self._shutdown(exc_type, exc_value, exc_tb)
+
+    def __call__[Function: typing.Callable[[], typing.AsyncGenerator[typing.Any, None]]](
+        self,
+        func: Function,
+        /,
+    ) -> Function:
+        self.lifespan_function = asynccontextmanager(func)
+        return func
+
+    @property
+    def started(self) -> bool:
+        return self._started
+
+    @staticmethod
+    async def _run_tasks(tasks: list[CoroutineTask[typing.Any]], /) -> None:
+        while tasks:
+            await tasks.pop(0)
+
+    async def _start(self) -> None:
+        if not self._started:
+            await logger.adebug("Running lifespan startup tasks")
+            self._started = True
+
+            if self.lifespan_function is not None:
+                self._lifespan_context = self.lifespan_function()
+                await self._lifespan_context.__aenter__()
+
+            await self._run_tasks(self.startup_tasks)
+
+    async def _shutdown(self, *suppress_args: typing.Any) -> None:
+        if self._started:
+            await logger.adebug("Running lifespan shutdown tasks")
+            self._started = False
+
+            if self._lifespan_context is not None:
+                await self._lifespan_context.__aexit__(*suppress_args)
+                self._lifespan_context = None
+
+            await self._run_tasks(self.shutdown_tasks)
+
+    def start(self, loop: asyncio.AbstractEventLoop | None = None) -> None:
+        run_task(self._start(), loop=loop)
+
+    def shutdown(self, loop: asyncio.AbstractEventLoop | None = None) -> None:
+        run_task(self._shutdown(None, None, None), loop=loop)
+
+    def on_startup[**P, T](self, task: Task[P, T], /) -> Task[P, T]:
+        self.startup_tasks.append(to_coroutine_task(task))
+        return task
+
+    def on_shutdown[**P, T](self, task: Task[P, T], /) -> Task[P, T]:
+        self.shutdown_tasks.append(to_coroutine_task(task))
+        return task
+
+
+__all__ = ("CoroutineTask", "DelayedTask", "Lifespan", "run_task", "to_coroutine_task")

telegrinder/tools/limited_dict.py

@@ -0,0 +1,35 @@
+from collections import UserDict, deque
+
+
+class LimitedDict[Key, Value](UserDict[Key, Value]):
+    def __init__(self, *, maxlimit: int = 1000) -> None:
+        super().__init__()
+        self.maxlimit = maxlimit
+        self.queue: deque[Key] = deque(maxlen=maxlimit)
+
+    def set(self, key: Key, value: Value, /) -> Value | None:
+        """Set item in the dictionary.
+        Returns a value that was deleted when the limit in the dictionary
+        was reached, otherwise None.
+        """
+        deleted_item = None
+
+        if len(self.queue) >= self.maxlimit:
+            deleted_item = self.pop(self.queue.popleft(), None)
+
+        if key not in self.queue:
+            self.queue.append(key)
+
+        super().__setitem__(key, value)
+        return deleted_item
+
+    def __setitem__(self, key: Key, value: Value, /) -> None:
+        self.set(key, value)
+
+    def __delitem__(self, key: Key, /) -> None:
+        if key in self.queue:
+            self.queue.remove(key)
+        return super().__delitem__(key)
+
+
+__all__ = ("LimitedDict",)

telegrinder/tools/loop_wrapper.py

@@ -0,0 +1,271 @@
+import asyncio
+import contextlib
+import datetime
+import enum
+import signal
+import typing
+from contextlib import suppress
+
+from telegrinder.modules import logger
+from telegrinder.tools.aio import TaskGroup, cancel_future, loop_is_running, run_task
+from telegrinder.tools.final import Final
+from telegrinder.tools.fullname import fullname
+from telegrinder.tools.lifespan import (
+    CoroutineFunc,
+    CoroutineTask,
+    DelayedTask,
+    Lifespan,
+    Task,
+    to_coroutine_task,
+)
+from telegrinder.tools.singleton.singleton import Singleton
+
+if typing.TYPE_CHECKING:
+    from contextvars import Context
+
+    type Tasks = set[asyncio.Task[typing.Any]]
+    type DelayedFunctionDecorator[**P, R] = typing.Callable[[typing.Callable[P, R]], DelayedFunction[P, R]]
+
+    class DelayedFunction[**P, R](typing.Protocol):
+        __name__: str
+        __delayed_task__: DelayedTask[typing.Callable[P, CoroutineTask[R]]]
+
+        def __call__(self, *args: P.args, **kwargs: P.kwargs) -> CoroutineTask[R]: ...
+
+        def cancel(self) -> bool:
+            """Cancel delayed task."""
+            ...
+
+
+class Timer(datetime.timedelta):
+    repeat = False
+
+    def __call__[**P, R](self, function: CoroutineFunc[P, R], /) -> DelayedFunction[P, R]:
+        loop_wrapper = LoopWrapper()
+        loop_wrapper.add_task(DelayedTask(function, seconds=self.total_seconds(), repeat=self.repeat))
+        return function  # type: ignore
+
+
+class Interval(Timer):
+    repeat = True
+
+
+@enum.unique
+class LoopWrapperState(enum.Enum):
+    NOT_RUNNING = enum.auto()
+    RUNNING = enum.auto()
+    RUNNING_MANUALLY = enum.auto()
+    SHUTDOWN = enum.auto()
+
+
+@typing.final
+class LoopWrapper(Singleton, Final):
+    _loop: asyncio.AbstractEventLoop
+    _lifespan: Lifespan
+    _future_tasks: list[CoroutineTask[typing.Any]]
+    _state: LoopWrapperState
+    _all_tasks: set[asyncio.Task[typing.Any]]
+    _event_stop: asyncio.Event
+    _run_lw_task: asyncio.Handle
+    _is_attached_to_running_loop: bool
+
+    timer = Timer
+    interval = Interval
+
+    __slots__ = (
+        "_lifespan",
+        "_future_tasks",
+        "_state",
+        "_all_tasks",
+        "_loop",
+        "_event_stop",
+        "_run_lw_task",
+        "_is_attached_to_running_loop",
+    )
+
+    def __init__(self) -> None:
+        try:
+            self._loop = asyncio.get_event_loop()
+        except RuntimeError:
+            self._loop = asyncio.new_event_loop()
+
+        self._run_lw_task = self._run_lw_later()
+        self._lifespan = Lifespan()
+        self._event_stop = asyncio.Event()
+        self._state = LoopWrapperState.NOT_RUNNING
+        self._future_tasks = [self._waiter_stop()]
+        self._all_tasks = set()
+        self._is_attached_to_running_loop = False
+
+        signal.signal(signal.SIGTERM, lambda *_: self.stop())
+
+    def __repr__(self) -> str:
+        return "<{}: loop={!r}, lifespan={!r}>".format(
+            fullname(self) + (" (running)" if self.running else ""),
+            self._loop,
+            self._lifespan,
+        )
+
+    def __call__(self) -> asyncio.AbstractEventLoop:
+        return self._loop
+
+    @property
+    def lifespan(self) -> Lifespan:
+        return self._lifespan
+
+    @property
+    def loop(self) -> asyncio.AbstractEventLoop:
+        return self._loop
+
+    @property
+    def time(self) -> float:
+        return self._loop.time()
+
+    @property
+    def running(self) -> bool:
+        return self._state in {LoopWrapperState.RUNNING, LoopWrapperState.RUNNING_MANUALLY}
+
+    @property
+    def shutdown(self) -> bool:
+        return self._state is LoopWrapperState.SHUTDOWN
+
+    async def _run_async_event_loop(self) -> None:
+        if not self.running:
+            self._state = LoopWrapperState.RUNNING
+            async with self._async_wrap_loop():
+                await self._run()
+
+    async def _run(self) -> None:
+        await logger.adebug("Running loop wrapper")
+
+        while self._future_tasks:
+            self._create_task(self._future_tasks.pop(0))
+
+        while self.running and (tasks := self._get_all_tasks()):
+            tasks_results, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+            for task_result in tasks_results:
+                try:
+                    task_result.result()
+                except Exception:
+                    await logger.aexception("Traceback message below:")
+
+    async def _cancel_tasks(self) -> None:
+        with suppress(asyncio.exceptions.CancelledError):
+            await cancel_future(asyncio.gather(*self._get_all_tasks(), return_exceptions=True))
+
+    @contextlib.asynccontextmanager
+    async def _async_wrap_loop(self) -> typing.AsyncGenerator[typing.Any, None]:
+        try:
+            try:
+                await self._lifespan._start()
+            finally:
+                yield
+        except asyncio.CancelledError:
+            await logger.adebug("Cancelling tasks...")
+            await self._cancel_tasks()
+        finally:
+            await self._shutdown()
+
+    async def _shutdown(self) -> None:
+        await self.lifespan._shutdown(None, None, None)
+        await logger.adebug("Shutting down loop wrapper")
+        self._state = LoopWrapperState.SHUTDOWN
+
+    async def _waiter_stop(self) -> None:
+        await self._event_stop.wait()
+        self._state = LoopWrapperState.SHUTDOWN
+        self._event_stop.clear()
+        await self._shutdown()
+        await self._cancel_tasks()
+
+    def _run_lw_later(self) -> asyncio.Handle:
+        return self._loop.call_soon_threadsafe(lambda l: l.create_task(self._run_async_event_loop()), self._loop)
+
+    def _create_task(
+        self,
+        coro: CoroutineTask[typing.Any],
+        /,
+        name: str | None = None,
+        context: Context | None = None,
+    ) -> asyncio.Task[typing.Any]:
+        task = self._loop.create_task(coro, name=name, context=context)
+        self._all_tasks.add(task)
+        task.add_done_callback(self._all_tasks.discard)
+
+        try:
+            return task
+        finally:
+            del task
+
+    def _get_all_tasks(self) -> Tasks:
+        """Get a set of all tasks from the loop wrapper and event loop (`exclude the current task if any`)."""
+
+        return (self._all_tasks | asyncio.all_tasks(loop=self._loop)).symmetric_difference(
+            set() if (task := asyncio.current_task(self._loop)) is None else {task},
+        )
+
+    def _close_loop(self) -> None:
+        if not self._loop.is_closed():
+            logger.debug("Closing event loop {!r}", self._loop)
+            self._loop.close()
+
+    @contextlib.contextmanager
+    def _wrap_loop(self, *, close_loop: bool = True) -> typing.Generator[typing.Any, None, None]:
+        try:
+            try:
+                self.lifespan.start(self._loop)
+            finally:
+                yield
+        except (KeyboardInterrupt, SystemExit) as e:
+            print(flush=True)
+            logger.info(f"Caught {e.__class__.__name__}, cancelling tasks")
+            run_task(self._cancel_tasks(), loop=self._loop)
+
+            if isinstance(e, SystemExit):
+                raise
+        finally:
+            run_task(self._shutdown(), loop=self._loop)
+
+            if close_loop:
+                self._close_loop()
+
+    def run(self, *, close_loop: bool = True) -> typing.NoReturn:  # type: ignore
+        if self.running:
+            raise RuntimeError("Loop wrapper already running.")
+
+        self._state = LoopWrapperState.RUNNING_MANUALLY
+        with self._wrap_loop(close_loop=close_loop):
+            run_task(self._run(), loop=self._loop)
+
+    def stop(self) -> None:
+        if not self._event_stop.is_set():
+            self._event_stop.set()
+
+    def add_task(self, task: Task[..., typing.Any], /) -> None:
+        coro_task = to_coroutine_task(task)
+
+        if not self.running:
+            self._future_tasks.append(coro_task)
+
+            if self._is_attached_to_running_loop is False and loop_is_running():
+                self.attach_to_running_loop()
+        else:
+            self._create_task(coro_task)
+
+    def attach_to_running_loop(self) -> None:
+        if self.running:
+            raise RuntimeError("Cannot attach the running loop wrapper to the running loop.")
+
+        self._loop = asyncio.get_running_loop()
+        self._run_lw_task.cancel()
+        self._is_attached_to_running_loop = True
+        self._run_lw_task = self._run_lw_later()
+
+    def create_task_group[T](self) -> TaskGroup[typing.Any]:
+        loop = asyncio.get_running_loop()
+        if not self.running and self._loop is not loop:
+            self.attach_to_running_loop()
+        return TaskGroup(loop)
+
+
+__all__ = ("DelayedTask", "LoopWrapper", "to_coroutine_task")

telegrinder/tools/magic/__init__.py

@@ -0,0 +1,29 @@
+from telegrinder.tools.magic.annotations import Annotations, get_generic_parameters
+from telegrinder.tools.magic.descriptors import additional_property
+from telegrinder.tools.magic.function import (
+    Bundle,
+    bundle,
+    get_default_args,
+    get_func_annotations,
+    get_func_parameters,
+    resolve_arg_names,
+    resolve_kwonly_arg_names,
+    resolve_posonly_arg_names,
+)
+from telegrinder.tools.magic.shortcut import Shortcut, shortcut
+
+__all__ = (
+    "Annotations",
+    "Bundle",
+    "Shortcut",
+    "additional_property",
+    "bundle",
+    "get_default_args",
+    "get_func_annotations",
+    "get_func_parameters",
+    "get_generic_parameters",
+    "resolve_arg_names",
+    "resolve_kwonly_arg_names",
+    "resolve_posonly_arg_names",
+    "shortcut",
+)

telegrinder/tools/magic/annotations.py

@@ -0,0 +1,172 @@
+import dataclasses
+import sys
+import types
+import typing
+from annotationlib import Format, ForwardRef, get_annotations
+from functools import cached_property
+from reprlib import recursive_repr
+
+from kungfu.library.misc import from_optional
+from kungfu.library.monad.option import NOTHING, Option, Some
+
+from telegrinder.tools.global_context.global_context import GlobalContext, ctx_var
+
+type TypeParameter = typing.Union[
+    typing.TypeVar,
+    typing.TypeVarTuple,
+    typing.ParamSpec,
+]
+type TypeParameters = tuple[TypeParameter, ...]
+type SupportsAnnotations = type[typing.Any] | types.ModuleType | typing.Callable[..., typing.Any]
+type AnnotationForm = typing.Any
+
+_UNION_TYPES: typing.Final = frozenset((typing.Union, types.UnionType))
+_CACHED_ANNOTATIONS: typing.Final = GlobalContext(
+    "cached_annotations",
+    thread_safe=True,
+    annotations=ctx_var(default_factory=dict, const=True),
+)
+
+
+def _cache_annotations(obj: SupportsAnnotations, annotations: dict[str, AnnotationForm], /) -> None:
+    _CACHED_ANNOTATIONS.annotations[obj] = MappingAnnotations(annotations)
+
+
+def _get_cached_annotations(obj: SupportsAnnotations, /) -> MappingAnnotations | None:
+    return _CACHED_ANNOTATIONS.annotations.get(obj)
+
+
+def is_union_type(obj: typing.Any, /) -> bool:
+    return typing.get_origin(obj) in _UNION_TYPES
+
+
+class MappingAnnotations[T = AnnotationForm](dict[str, T]):
+    def __init__(self, annotations: typing.Mapping[str, AnnotationForm], /) -> None:
+        super().__init__(annotations)
+        self.return_type: Option[T] = from_optional(self.pop("return", None))
+
+    @recursive_repr()
+    def __repr__(self) -> str:
+        return f"annotations={super().__repr__()}, return_type={self.return_type!r}"
+
+
+@dataclasses.dataclass
+class Annotations:
+    obj: SupportsAnnotations
+
+    @cached_property
+    def forward_ref_parameters(self) -> dict[str, typing.Any]:
+        parameters = dict[str, typing.Any](
+            is_argument=False,
+            is_class=False,
+            module=None,
+        )
+
+        if isinstance(self.obj, type):
+            parameters["is_class"] = True
+            parameters["module"] = (
+                sys.modules[module] if (module := getattr(self.obj, "__module__", None)) is not None else None
+            )
+        elif isinstance(self.obj, types.ModuleType):
+            parameters["module"] = self.obj
+        elif callable(self.obj):
+            parameters["is_argument"] = True
+
+        return parameters
+
+    @cached_property
+    def generic_parameters(self) -> Option[dict[TypeParameter, typing.Any]]:
+        return get_generic_parameters(self.obj)
+
+    @classmethod
+    def from_obj(cls, obj: typing.Any, /) -> typing.Self:
+        if not isinstance(obj, type | types.ModuleType | typing.Callable):
+            obj = type(obj)
+
+        return cls(obj)
+
+    @typing.overload
+    def get(
+        self,
+        *,
+        ignore_failed_evals: bool = True,
+        cache: bool = False,
+    ) -> MappingAnnotations[typing.ForwardRef | AnnotationForm]: ...
+
+    @typing.overload
+    def get(
+        self,
+        *,
+        exclude_forward_refs: typing.Literal[True],
+        ignore_failed_evals: bool = True,
+        cache: bool = False,
+    ) -> MappingAnnotations: ...
+
+    def get(
+        self,
+        *,
+        exclude_forward_refs: bool = False,
+        ignore_failed_evals: bool = True,
+        cache: bool = False,
+    ) -> MappingAnnotations:
+        if (cached_annotations := _get_cached_annotations(self.obj)) is not None:
+            return cached_annotations
+
+        annotations = dict[str, typing.Any]()
+        for name, annotation in get_annotations(
+            obj=self.obj,
+            format=Format.FORWARDREF,
+            eval_str=False,
+        ).items():
+            if isinstance(annotation, str):
+                annotation = ForwardRef(annotation, owner=self.obj, **self.forward_ref_parameters)
+
+            if not isinstance(annotation, ForwardRef):
+                annotations[name] = annotation
+                continue
+
+            try:
+                value = annotation.evaluate(format=Format.VALUE)
+            except NameError:
+                if not ignore_failed_evals:
+                    raise
+
+                value = annotation
+
+            if isinstance(value, typing.ForwardRef) and exclude_forward_refs:
+                continue
+
+            annotations[name] = value
+
+        if cache:
+            _cache_annotations(self.obj, annotations)
+
+        return MappingAnnotations(annotations)
+
+
+def get_generic_parameters(obj: typing.Any, /) -> Option[dict[TypeParameter, AnnotationForm]]:
+    origin_obj = typing.get_origin(obj)
+    args = typing.get_args(obj)
+    parameters: TypeParameters = getattr(origin_obj or obj, "__parameters__")
+
+    if not parameters:
+        return NOTHING
+
+    index = 0
+    generic_alias_args = dict[TypeParameter, typing.Any]()
+
+    for parameter in parameters:
+        if isinstance(parameter, typing.TypeVarTuple):
+            stop_index = len(args) - index
+            generic_alias_args[parameter] = args[index:stop_index]
+            index = stop_index
+            continue
+
+        arg = args[index] if index < len(args) else None
+        generic_alias_args[parameter] = arg
+        index += 1
+
+    return Some(generic_alias_args)
+
+
+__all__ = ("Annotations", "get_generic_parameters")