telegrinder 0.3.4__py3-none-any.whl → 0.4.0__py3-none-any.whl

telegrinder/tools/lifespan.py

@@ -0,0 +1,105 @@
+import asyncio
+import dataclasses
+import datetime
+import typing
+
+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[typing.Callable[P, CoroutineTask[T]]]
+
+
+def run_tasks(
+    tasks: list[CoroutineTask[typing.Any]],
+    /,
+) -> None:
+    loop = asyncio.get_event_loop()
+    while tasks:
+        loop.run_until_complete(tasks.pop(0))
+
+
+def to_coroutine_task[**P, T](task: Task[P, 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(slots=True)
+class DelayedTask[Handler: CoroutineFunc[..., typing.Any]]:
+    handler: Handler
+    seconds: float | datetime.timedelta
+    repeat: bool = dataclasses.field(default=False, kw_only=True)
+    _cancelled: bool = dataclasses.field(default=False, init=False, repr=False)
+
+    @property
+    def is_cancelled(self) -> bool:
+        return self._cancelled
+
+    @property
+    def delay(self) -> float:
+        return self.seconds if isinstance(self.seconds, int | float) else self.seconds.total_seconds()
+
+    async def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> typing.Any:
+        while not self.is_cancelled:
+            await asyncio.sleep(self.delay)
+            if self.is_cancelled:
+                break
+            try:
+                await self.handler(*args, **kwargs)
+            finally:
+                if not self.repeat:
+                    break
+
+    def cancel(self) -> None:
+        if not self._cancelled:
+            self._cancelled = True
+
+
+@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
+class Lifespan:
+    startup_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
+    shutdown_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
+
+    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
+
+    def start(self) -> None:
+        run_tasks(self.startup_tasks)
+
+    def shutdown(self) -> None:
+        run_tasks(self.shutdown_tasks)
+
+    def __enter__(self) -> None:
+        self.start()
+
+    def __exit__(self) -> None:
+        self.shutdown()
+
+    async def __aenter__(self) -> None:
+        for task in self.startup_tasks:
+            await task
+
+    async def __aexit__(self, *args) -> None:
+        for task in self.shutdown_tasks:
+            await task
+
+    def __add__(self, other: typing.Self, /) -> typing.Self:
+        return self.__class__(
+            startup_tasks=self.startup_tasks + other.startup_tasks,
+            shutdown_tasks=self.shutdown_tasks + other.shutdown_tasks,
+        )
+
+
+__all__ = (
+    "CoroutineTask",
+    "DelayedTask",
+    "Lifespan",
+    "run_tasks",
+    "to_coroutine_task",
+)

telegrinder/tools/limited_dict.py

@@ -1,37 +1,32 @@
-import typing
-from collections import UserDict, deque
-
-KT = typing.TypeVar("KT")
-VT = typing.TypeVar("VT")
-
-
-class LimitedDict(UserDict[KT, VT]):
-    def __init__(self, *, maxlimit: int = 1000) -> None:
-        super().__init__()
-        self.maxlimit = maxlimit
-        self.queue: deque[KT] = deque(maxlen=maxlimit)
-
-    def set(self, key: KT, value: VT, /) -> VT | 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: KT, value: VT, /) -> None:
-        self.set(key, value)
-
-    def __delitem__(self, key: KT) -> None:
-        if key in self.queue:
-            self.queue.remove(key)
-        return super().__delitem__(key)
-
-
-__all__ = ("LimitedDict",)
+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/__init__.py

@@ -1,4 +1,4 @@
-from .abc import ABCLoopWrapper
-from .loop_wrapper import DelayedTask, Lifespan, LoopWrapper
-
-__all__ = ("ABCLoopWrapper", "DelayedTask", "Lifespan", "LoopWrapper")
+from .abc import ABCLoopWrapper
+from .loop_wrapper import DelayedTask, Lifespan, LoopWrapper
+
+__all__ = ("ABCLoopWrapper", "DelayedTask", "Lifespan", "LoopWrapper")

telegrinder/tools/loop_wrapper/abc.py

@@ -1,15 +1,20 @@
-import typing
-from abc import ABC, abstractmethod
-
-
-class ABCLoopWrapper(ABC):
-    @abstractmethod
-    def add_task(self, task: typing.Any) -> None:
-        pass
-
-    @abstractmethod
-    def run_event_loop(self) -> None:
-        pass
-
-
-__all__ = ("ABCLoopWrapper",)
+import typing
+from abc import ABC, abstractmethod
+
+
+class ABCLoopWrapper(ABC):
+    @property
+    @abstractmethod
+    def is_running(self) -> bool:
+        pass
+
+    @abstractmethod
+    def add_task(self, task: typing.Any, /) -> None:
+        pass
+
+    @abstractmethod
+    def run_event_loop(self) -> typing.NoReturn:
+        raise NotImplementedError
+
+
+__all__ = ("ABCLoopWrapper",)

telegrinder/tools/loop_wrapper/loop_wrapper.py

@@ -1,224 +1,169 @@
-import asyncio
-import contextlib
-import dataclasses
-import datetime
-import typing
-
-from telegrinder.modules import logger
-from telegrinder.tools.loop_wrapper.abc import ABCLoopWrapper
-
-T = typing.TypeVar("T")
-P = typing.ParamSpec("P")
-CoroFunc = typing.TypeVar("CoroFunc", bound="CoroutineFunc[..., typing.Any]")
-
-CoroutineTask: typing.TypeAlias = typing.Coroutine[typing.Any, typing.Any, T]
-CoroutineFunc: typing.TypeAlias = typing.Callable[P, CoroutineTask[T]]
-Task: typing.TypeAlias = (
-    "CoroutineFunc[P, T] | CoroutineTask[T] | DelayedTask[typing.Callable[P, CoroutineTask[T]]]"
-)
-
-
-def run_tasks(
-    tasks: list[CoroutineTask[typing.Any]],
-    loop: asyncio.AbstractEventLoop,
-) -> None:
-    while tasks:
-        loop.run_until_complete(tasks.pop(0))
-
-
-def to_coroutine_task(task: Task) -> CoroutineTask[typing.Any]:
-    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(slots=True)
-class DelayedTask(typing.Generic[CoroFunc]):
-    handler: CoroFunc
-    seconds: float
-    repeat: bool = dataclasses.field(default=False, kw_only=True)
-    _cancelled: bool = dataclasses.field(default=False, init=False, repr=False)
-
-    @property
-    def is_cancelled(self) -> bool:
-        return self._cancelled
-
-    async def __call__(self, *args: typing.Any, **kwargs: typing.Any) -> None:
-        while not self.is_cancelled:
-            await asyncio.sleep(self.seconds)
-            if self.is_cancelled:
-                break
-            try:
-                await self.handler(*args, **kwargs)
-            finally:
-                if not self.repeat:
-                    break
-
-    def cancel(self) -> None:
-        if not self._cancelled:
-            self._cancelled = True
-
-
-@dataclasses.dataclass(kw_only=True, slots=True, frozen=True)
-class Lifespan:
-    startup_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
-    shutdown_tasks: list[CoroutineTask[typing.Any]] = dataclasses.field(default_factory=lambda: [])
-
-    def on_startup(self, task: Task, /) -> Task:
-        self.startup_tasks.append(to_coroutine_task(task))
-        return task
-
-    def on_shutdown(self, task: Task, /) -> Task:
-        self.shutdown_tasks.append(to_coroutine_task(task))
-        return task
-
-    def start(self, loop: asyncio.AbstractEventLoop, /) -> None:
-        run_tasks(self.startup_tasks, loop)
-
-    def shutdown(self, loop: asyncio.AbstractEventLoop, /) -> None:
-        run_tasks(self.shutdown_tasks, loop)
-
-
-class LoopWrapper(ABCLoopWrapper):
-    def __init__(
-        self,
-        *,
-        tasks: list[CoroutineTask[typing.Any]] | None = None,
-        lifespan: Lifespan | None = None,
-        event_loop: asyncio.AbstractEventLoop | None = None,
-    ) -> None:
-        self.tasks: list[CoroutineTask[typing.Any]] = tasks or []
-        self.lifespan = lifespan or Lifespan()
-        self._loop = event_loop
-
-    @property
-    def loop(self) -> asyncio.AbstractEventLoop:
-        assert self._loop is not None
-        return self._loop
-
-    def __repr__(self) -> str:
-        return "<{}: loop={!r} with tasks={!r}, lifespan={!r}>".format(
-            self.__class__.__name__,
-            self._loop,
-            self.tasks,
-            self.lifespan,
-        )
-
-    def run_event_loop(self) -> None:
-        if not self.tasks:
-            logger.warning("You run loop with 0 tasks!")
-
-        self._loop = asyncio.new_event_loop() if self._loop is None else self._loop
-        self.lifespan.start(self._loop)
-
-        while self.tasks:
-            self._loop.create_task(self.tasks.pop(0))
-
-        tasks = asyncio.all_tasks(self._loop)
-        try:
-            while tasks:
-                tasks_results, _ = self._loop.run_until_complete(
-                    asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION),
-                )
-                for task_result in tasks_results:
-                    try:
-                        task_result.result()
-                    except BaseException as ex:
-                        logger.exception(ex)
-                tasks = asyncio.all_tasks(self._loop)
-        except KeyboardInterrupt:
-            print()  # blank print for ^C
-            logger.info("Caught KeyboardInterrupt, cancellation...")
-            self.complete_tasks(tasks)
-        finally:
-            self.lifespan.shutdown(self._loop)
-            if self._loop.is_running():
-                self._loop.close()
-
-    def add_task(self, task: Task) -> None:
-        task = to_coroutine_task(task)
-
-        if self._loop is not None and self._loop.is_running():
-            self._loop.create_task(task)
-        else:
-            self.tasks.append(task)
-
-    def complete_tasks(self, tasks: set[asyncio.Task[typing.Any]]) -> None:
-        tasks = tasks | asyncio.all_tasks(self.loop)
-        task_to_cancel = asyncio.gather(*tasks, return_exceptions=True)
-        task_to_cancel.cancel()
-        with contextlib.suppress(asyncio.CancelledError):
-            self.loop.run_until_complete(task_to_cancel)
-
-    @typing.overload
-    def timer(self, *, seconds: datetime.timedelta) -> typing.Callable[..., typing.Any]: ...
-
-    @typing.overload
-    def timer(
-        self,
-        *,
-        days: int = 0,
-        hours: int = 0,
-        minutes: int = 0,
-        seconds: float = 0,
-    ) -> typing.Callable[..., typing.Any]: ...
-
-    def timer(
-        self,
-        *,
-        days: int = 0,
-        hours: int = 0,
-        minutes: int = 0,
-        seconds: float | datetime.timedelta = 0,
-    ) -> typing.Callable[..., typing.Any]:
-        if isinstance(seconds, datetime.timedelta):
-            seconds = seconds.total_seconds()
-
-        seconds += minutes * 60
-        seconds += hours * 60 * 60
-        seconds += days * 24 * 60 * 60
-
-        def decorator(func: CoroFunc) -> CoroFunc:
-            self.add_task(DelayedTask(func, seconds, repeat=False))
-            return func
-
-        return decorator
-
-    @typing.overload
-    def interval(self, *, seconds: datetime.timedelta) -> typing.Callable[..., typing.Any]: ...
-
-    @typing.overload
-    def interval(
-        self,
-        *,
-        days: int = 0,
-        hours: int = 0,
-        minutes: int = 0,
-        seconds: float = 0,
-    ) -> typing.Callable[..., typing.Any]: ...
-
-    def interval(
-        self,
-        *,
-        days: int = 0,
-        hours: int = 0,
-        minutes: int = 0,
-        seconds: float | datetime.timedelta = 0,
-    ) -> typing.Callable[..., typing.Any]:
-        if isinstance(seconds, datetime.timedelta):
-            seconds = seconds.total_seconds()
-
-        seconds += minutes * 60
-        seconds += hours * 60 * 60
-        seconds += days * 24 * 60 * 60
-
-        def decorator(func: CoroFunc) -> CoroFunc:
-            self.add_task(DelayedTask(func, seconds, repeat=True))
-            return func
-
-        return decorator
-
-
-__all__ = ("DelayedTask", "Lifespan", "LoopWrapper", "to_coroutine_task")
+import asyncio
+import contextlib
+import datetime
+import typing
+
+from telegrinder.modules import logger
+from telegrinder.tools.lifespan import (
+    CoroutineFunc,
+    CoroutineTask,
+    DelayedTask,
+    Lifespan,
+    Task,
+    to_coroutine_task,
+)
+from telegrinder.tools.loop_wrapper.abc import ABCLoopWrapper
+from telegrinder.tools.magic import cancel_future
+
+
+class LoopWrapper(ABCLoopWrapper):
+    def __init__(
+        self,
+        *,
+        tasks: list[CoroutineTask[typing.Any]] | None = None,
+        lifespan: Lifespan | None = None,
+    ) -> None:
+        self.tasks: list[CoroutineTask[typing.Any]] = tasks or []
+        self.lifespan = lifespan or Lifespan()
+        self._loop: asyncio.AbstractEventLoop | None = None
+
+    @property
+    def is_running(self) -> bool:
+        if self._loop is None:
+            return False
+        return self._loop.is_running()
+
+    @property
+    def loop(self) -> asyncio.AbstractEventLoop:
+        assert self._loop is not None, "Loop is not set."
+        return self._loop
+
+    def __repr__(self) -> str:
+        return "<{}: loop={!r} with tasks={!r}, lifespan={!r}>".format(
+            self.__class__.__name__,
+            self._loop,
+            self.tasks,
+            self.lifespan,
+        )
+
+    async def _run_tasks(self) -> None:
+        async with asyncio.TaskGroup() as tg:
+            while self.tasks:
+                tg.create_task(self.tasks.pop(0))
+
+    def run_event_loop(self) -> typing.NoReturn:  # type: ignore
+        if not self.tasks:
+            logger.warning("Run loop without tasks!")
+
+        try:
+            self._loop = asyncio.get_running_loop()
+        except RuntimeError:
+            self._loop = asyncio.get_event_loop()
+
+        self.lifespan.start()
+        self._loop.create_task(self._run_tasks())
+
+        tasks = asyncio.all_tasks(self._loop)
+        try:
+            while tasks:
+                tasks_results, _ = self._loop.run_until_complete(
+                    asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION),
+                )
+                for task_result in tasks_results:
+                    try:
+                        task_result.result()
+                    except BaseException:
+                        logger.exception("Traceback message below:")
+                tasks = asyncio.all_tasks(self._loop)
+        except KeyboardInterrupt:
+            print()  # blank print for ^C
+            logger.info("Caught KeyboardInterrupt, cancellation...")
+            self.complete_tasks(tasks)
+        finally:
+            self.lifespan.shutdown()
+            if self._loop.is_running():
+                self._loop.close()
+
+    def add_task(self, task: Task[..., typing.Any], /) -> None:
+        task = to_coroutine_task(task)
+
+        if self._loop is not None and self._loop.is_running():
+            self._loop.create_task(task)
+        else:
+            self.tasks.append(task)
+
+    def complete_tasks(self, tasks: set[asyncio.Task[typing.Any]], /) -> None:
+        tasks = tasks | asyncio.all_tasks(self.loop)
+        with contextlib.suppress(asyncio.CancelledError, asyncio.InvalidStateError):
+            self.loop.run_until_complete(cancel_future(asyncio.gather(*tasks, return_exceptions=True)))
+
+    @typing.overload
+    def timer(self, *, seconds: datetime.timedelta) -> typing.Callable[..., typing.Any]: ...
+
+    @typing.overload
+    def timer(
+        self,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float = 0,
+    ) -> typing.Callable[..., typing.Any]: ...
+
+    def timer(
+        self,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float | datetime.timedelta = 0,
+    ) -> typing.Callable[..., typing.Any]:
+        if isinstance(seconds, datetime.timedelta):
+            seconds = seconds.total_seconds()
+
+        seconds += minutes * 60
+        seconds += hours * 60 * 60
+        seconds += days * 24 * 60 * 60
+
+        def decorator[Func: CoroutineFunc[..., typing.Any]](func: Func) -> Func:
+            self.add_task(DelayedTask(func, seconds, repeat=False))
+            return func
+
+        return decorator
+
+    @typing.overload
+    def interval(self, *, seconds: datetime.timedelta) -> typing.Callable[..., typing.Any]: ...
+
+    @typing.overload
+    def interval(
+        self,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float = 0,
+    ) -> typing.Callable[..., typing.Any]: ...
+
+    def interval(
+        self,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float | datetime.timedelta = 0,
+    ) -> typing.Callable[..., typing.Any]:
+        if isinstance(seconds, datetime.timedelta):
+            seconds = seconds.total_seconds()
+
+        seconds += minutes * 60
+        seconds += hours * 60 * 60
+        seconds += days * 24 * 60 * 60
+
+        def decorator[Func: CoroutineFunc[..., typing.Any]](func: Func) -> Func:
+            self.add_task(DelayedTask(func, seconds, repeat=True))
+            return func
+
+        return decorator
+
+
+__all__ = ("DelayedTask", "LoopWrapper", "to_coroutine_task")