telegrinder 0.4.2__py3-none-any.whl → 0.5.1__py3-none-any.whl

telegrinder/tools/loop_wrapper.py

@@ -0,0 +1,332 @@
+import asyncio
+import contextlib
+import datetime
+import enum
+import typing
+
+from telegrinder.modules import logger
+from telegrinder.tools.aio import cancel_future, 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
+
+type Tasks = set[asyncio.Task[typing.Any]]
+type LoopFactory = typing.Callable[[], asyncio.AbstractEventLoop]
+type DelayedFunctionDecorator[**P, R] = typing.Callable[[typing.Callable[P, R]], DelayedFunction[P, R]]
+
+
+class DelayedFunction[**P, R](typing.Protocol):
+    __name__: str
+    __delayed_task__: DelayedTask
+
+    def __call__(self, *args: P.args, **kwargs: P.kwargs) -> typing.Coroutine[typing.Any, typing.Any, R]: ...
+
+    def cancel(self) -> bool:
+        """Cancel delayed task."""
+        ...
+
+
+@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
+    _tasks: list[CoroutineTask[typing.Any]]
+    _state: LoopWrapperState
+    _all_tasks: set[asyncio.Task[typing.Any]]
+    _limit: int | None
+    _semaphore: asyncio.Semaphore | None
+
+    __slots__ = (
+        "_loop",
+        "_lifespan",
+        "_tasks",
+        "_state",
+        "_all_tasks",
+        "_limit",
+        "_semaphore",
+    )
+
+    def __init__(self) -> None:
+        self._loop = asyncio.get_event_loop()
+        self._lifespan = Lifespan()
+        self._tasks = list()
+        self._state = LoopWrapperState.NOT_RUNNING
+        self._all_tasks = set()
+        self._limit = None
+        self._semaphore = None
+
+        self._create_task(self._run_async_event_loop())
+
+    def __call__(self) -> asyncio.AbstractEventLoop:
+        """A loop factory."""
+        return self._loop
+
+    def __repr__(self) -> str:
+        return "<{}: loop={!r} lifespan={!r}>".format(
+            fullname(self) + (" (running)" if self.running else ""),
+            self._loop,
+            self._lifespan,
+        )
+
+    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:
+        logger.debug("Running loop wrapper")
+
+        while self._tasks:
+            await self.create_task(self._tasks.pop(0))
+
+        while self.running and (tasks := self._get_all_tasks()):
+            await self._process_tasks(tasks)
+
+    async def _process_tasks(self, tasks: Tasks, /) -> None:
+        """Processes the given tasks, checks for exceptions, and raises them if any."""
+
+        tasks_results, _ = await asyncio.wait(tasks, return_when=asyncio.FIRST_EXCEPTION)
+        for task_result in tasks_results:
+            try:
+                if not task_result.cancelled() and (exception := task_result.exception()) is not None:
+                    raise exception from None  # Raise the exception that was set on the task.
+            except BaseException:
+                logger.exception("Traceback message below:")
+
+    async def _cancel_tasks(self) -> None:
+        with contextlib.suppress(asyncio.CancelledError, asyncio.InvalidStateError):
+            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:
+            await self._lifespan._start()
+            yield
+        except asyncio.CancelledError:
+            logger.info("LoopWrapper task was cancelled, cancelling tasks...")
+            await self._cancel_tasks()
+        finally:
+            await self._shutdown()
+
+    async def _shutdown(self) -> None:
+        await self.lifespan._shutdown()
+        logger.debug("Shutting down loop wrapper")
+        self._state = LoopWrapperState.SHUTDOWN
+
+    async def _run_coro_with_semaphore(self, coro: CoroutineTask[typing.Any], /) -> None:
+        assert self._semaphore is not None
+        try:
+            await coro
+        finally:
+            self._semaphore.release()
+
+    async def _create_task_with_semaphore(self, coro: CoroutineTask[typing.Any], /) -> None:
+        assert self._semaphore is not None
+        await self._semaphore.acquire()
+        self._create_task(self._run_coro_with_semaphore(coro))
+
+    def _create_task(self, coro: CoroutineTask[typing.Any], /) -> None:
+        task = self._loop.create_task(coro)
+        self._all_tasks.add(task)
+        task.add_done_callback(self._all_tasks.discard)
+
+    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:
+            self.lifespan.start()
+            yield
+        except KeyboardInterrupt:
+            logger.info("Caught KeyboardInterrupt, cancelling tasks...")
+            run_task(self._cancel_tasks())
+        finally:
+            run_task(self._shutdown())
+            if close_loop:
+                self._close_loop()
+
+    def _get_delayed_task_decorator(
+        self,
+        repeat: bool,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float | datetime.timedelta = 0.0,
+    ) -> typing.Callable[..., typing.Any]:
+        if isinstance(seconds, int | float):
+            seconds += minutes * 60
+            seconds += hours * 60 * 60
+            seconds += days * 24 * 60 * 60
+
+        def decorator[Func: CoroutineFunc[..., typing.Any]](function: Func) -> Func:
+            self.add_task(DelayedTask(function, seconds, repeat=repeat))
+            return function
+
+        return decorator
+
+    @property
+    def lifespan(self) -> Lifespan:
+        return self._lifespan
+
+    @property
+    def loop(self) -> asyncio.AbstractEventLoop:
+        return self._loop
+
+    @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
+
+    @property
+    def tasks_limit(self) -> int | None:
+        return self._limit
+
+    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())
+
+    @typing.overload
+    def bind_loop(self, *, loop_factory: LoopFactory) -> typing.Self: ...
+
+    @typing.overload
+    def bind_loop(self, *, loop: asyncio.AbstractEventLoop) -> typing.Self: ...
+
+    def bind_loop(
+        self,
+        *,
+        loop_factory: LoopFactory | None = None,
+        loop: asyncio.AbstractEventLoop | None = None,
+    ) -> typing.Self:
+        assert loop is not None or loop_factory is not None
+
+        if self.running:
+            logger.warning("Cannot bind a new event loop to running loop wrapper.")
+            return self
+
+        old_loop = self._loop
+        self._loop = loop_factory() if loop_factory else loop or self._loop
+
+        if old_loop is not self._loop:
+            self._create_task(self._run_async_event_loop())
+
+        return self
+
+    def limit(self, value: int, /) -> typing.Self:
+        if self._limit is not None:
+            raise ValueError("Cannot reset limit value.")
+
+        self._limit = value
+        self._semaphore = asyncio.Semaphore(value)
+        return self
+
+    def add_task(self, task: Task[..., typing.Any], /) -> None:
+        coro_task = to_coroutine_task(task)
+        return self._create_task(coro_task) if self.running else self._tasks.append(coro_task)
+
+    async def create_task(self, task: Task[..., typing.Any], /) -> None:
+        if not self.running:
+            self.add_task(task)
+            return
+
+        coro_task = to_coroutine_task(task)
+        if self._semaphore is not None:
+            await self._create_task_with_semaphore(coro_task)
+        else:
+            self._create_task(coro_task)
+
+    @typing.overload
+    def timer[**P, R](self, delta: datetime.timedelta, /) -> DelayedFunctionDecorator[P, R]: ...
+
+    @typing.overload
+    def timer[**P, R](
+        self,
+        *,
+        days: int = ...,
+        hours: int = ...,
+        minutes: int = ...,
+        seconds: float = ...,
+    ) -> DelayedFunctionDecorator[P, R]: ...
+
+    def timer(
+        self,
+        delta: datetime.timedelta | None = None,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float = 0.0,
+    ) -> DelayedFunctionDecorator[..., typing.Any]:
+        return self._get_delayed_task_decorator(
+            repeat=False,
+            days=days,
+            hours=hours,
+            minutes=minutes,
+            seconds=delta or seconds,
+        )
+
+    @typing.overload
+    def interval[**P, R](self, delta: datetime.timedelta, /) -> DelayedFunctionDecorator[P, R]: ...
+
+    @typing.overload
+    def interval[**P, R](
+        self,
+        *,
+        days: int = ...,
+        hours: int = ...,
+        minutes: int = ...,
+        seconds: float = ...,
+    ) -> DelayedFunctionDecorator[P, R]: ...
+
+    def interval(
+        self,
+        delta: datetime.timedelta | None = None,
+        *,
+        days: int = 0,
+        hours: int = 0,
+        minutes: int = 0,
+        seconds: float = 0.0,
+    ) -> DelayedFunctionDecorator[..., typing.Any]:
+        return self._get_delayed_task_decorator(
+            repeat=True,
+            days=days,
+            hours=hours,
+            minutes=minutes,
+            seconds=delta or seconds,
+        )
+
+
+__all__ = ("DelayedTask", "LoopWrapper", "to_coroutine_task")

telegrinder/tools/magic/__init__.py

@@ -0,0 +1,32 @@
+from telegrinder.tools.magic.annotations import Annotations, get_generic_parameters
+from telegrinder.tools.magic.dictionary import extract, join_dicts
+from telegrinder.tools.magic.function import (
+    Bundle,
+    bundle,
+    function_context,
+    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",
+    "bundle",
+    "extract",
+    "function_context",
+    "get_default_args",
+    "get_func_annotations",
+    "get_func_parameters",
+    "get_generic_parameters",
+    "join_dicts",
+    "resolve_arg_names",
+    "resolve_kwonly_arg_names",
+    "resolve_posonly_arg_names",
+    "shortcut",
+)

telegrinder/tools/magic/annotations.py

@@ -0,0 +1,165 @@
+import dataclasses
+import sys
+import types
+import typing as _typing
+from functools import cached_property
+
+import typing_extensions as typing
+from fntypes.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,
+    _typing.TypeVar,
+    _typing.TypeVarTuple,
+    _typing.ParamSpec,
+]
+type TypeParameters = tuple[TypeParameter, ...]
+type SupportsAnnotations = type[typing.Any] | types.ModuleType | typing.Callable[..., typing.Any]
+
+_CACHED_ANNOTATIONS: typing.Final[GlobalContext] = GlobalContext(
+    "cached_annotations",
+    annotations=ctx_var(default_factory=dict, const=True),
+)
+
+
+def _cache_annotations(obj: SupportsAnnotations, annotations: dict[str, typing.Any], /) -> None:
+    _CACHED_ANNOTATIONS.annotations[obj] = annotations
+
+
+def _get_cached_annotations(obj: SupportsAnnotations, /) -> dict[str, typing.Any] | None:
+    return _CACHED_ANNOTATIONS.annotations.get(obj)
+
+
+@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,
+        allow_return_type: bool = False,
+        cache: bool = False,
+    ) -> dict[str, typing.Any | typing.ForwardRef]: ...
+
+    @typing.overload
+    def get(
+        self,
+        *,
+        exclude_forward_refs: typing.Literal[True],
+        ignore_failed_evals: bool = True,
+        allow_return_type: bool = False,
+        cache: bool = False,
+    ) -> dict[str, typing.Any]: ...
+
+    def get(
+        self,
+        *,
+        exclude_forward_refs: bool = False,
+        ignore_failed_evals: bool = True,
+        allow_return_type: bool = True,
+        cache: bool = False,
+    ) -> dict[str, typing.Any]:
+        if (cached_annotations := _get_cached_annotations(self.obj)) is not None:
+            return cached_annotations
+
+        annotations = dict[str, typing.Any]()
+        for name, annotation in typing.get_annotations(
+            obj=self.obj,
+            format=typing.Format.FORWARDREF,
+        ).items():
+            if isinstance(annotation, str):
+                annotation = typing.ForwardRef(annotation, **self.forward_ref_parameters)
+
+            if not isinstance(annotation, typing.ForwardRef):
+                annotations[name] = annotation
+                continue
+
+            try:
+                value = typing.evaluate_forward_ref(
+                    forward_ref=annotation,
+                    owner=self.obj,
+                    format=typing.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 not allow_return_type:
+            annotations.pop("return", None)
+
+        if cache:
+            _cache_annotations(self.obj, annotations)
+
+        return annotations
+
+
+def get_generic_parameters(obj: typing.Any, /) -> Option[dict[TypeParameter, typing.Any]]:
+    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")

telegrinder/tools/magic/dictionary.py

@@ -0,0 +1,20 @@
+import typing
+
+
+def join_dicts[Key: typing.Hashable, Value](
+    left_dict: dict[Key, typing.Any],
+    right_dict: dict[typing.Any, Value],
+    /,
+) -> dict[Key, Value]:
+    return {left_key: right_dict[right_key] for left_key, right_key in left_dict.items()}
+
+
+def extract[Key: typing.Hashable, Value](
+    keys: typing.Iterable[Key],
+    mapping: typing.Mapping[Key, Value],
+    /,
+) -> dict[Key, Value]:
+    return {key: mapping[key] for key in keys if key in mapping}
+
+
+__all__ = ("extract", "join_dicts")