telegrinder 1.0.0rc1__py3-none-any.whl

telegrinder/bot/dispatch/waiter_machine/machine.py

@@ -0,0 +1,264 @@
+import asyncio
+import datetime
+import typing
+
+from telegrinder.bot.cute_types.base import BaseCute
+from telegrinder.bot.dispatch.context import Context
+from telegrinder.bot.dispatch.waiter_machine.actions import WaiterActions
+from telegrinder.bot.dispatch.waiter_machine.hasher import Hasher
+from telegrinder.bot.dispatch.waiter_machine.middleware import WaiterMiddleware
+from telegrinder.bot.dispatch.waiter_machine.short_state import (
+    ShortState,
+    ShortStateContext,
+)
+from telegrinder.bot.rules.abc import ABCRule
+from telegrinder.modules import logger
+from telegrinder.tools.aio import maybe_awaitable
+from telegrinder.tools.lifespan import Lifespan
+from telegrinder.tools.limited_dict import LimitedDict
+from telegrinder.tools.magic.function import bundle
+
+if typing.TYPE_CHECKING:
+    from telegrinder.bot.dispatch.view.base import View
+
+type Storage[Event: BaseCute, HasherData] = dict[
+    Hasher[Event, HasherData],
+    LimitedDict[typing.Hashable, ShortState[Event]],
+]
+type HasherWithData[Event: BaseCute, ViewType: View, Data] = tuple[Hasher[Event, Data], ViewType, Data]
+
+_NODATA: typing.Final = object()
+MAX_STORAGE_SIZE: typing.Final = 10000
+ONE_MINUTE: typing.Final = datetime.timedelta(minutes=1)
+WEEK: typing.Final = datetime.timedelta(days=7)
+
+
+def unpack_to_context(context: Context, /) -> tuple[Context]:
+    return (context,)
+
+
+def no_unpack(_: Context, /) -> tuple[()]:
+    return ()
+
+
+class ContextUnpackProto[*Ts](typing.Protocol):
+    __name__: str
+
+    def __call__(self, context: Context, /) -> tuple[*Ts]: ...
+
+
+class WaiterMachine:
+    storage: Storage[typing.Any, typing.Any]
+
+    def __init__(
+        self,
+        *,
+        max_storage_size: int = MAX_STORAGE_SIZE,
+        base_state_lifetime: datetime.timedelta = WEEK,
+        clear_storage_interval: datetime.timedelta = ONE_MINUTE,
+    ) -> None:
+        self.max_storage_size = max_storage_size
+        self.base_state_lifetime = base_state_lifetime
+        self.storage = {}
+
+    def __repr__(self) -> str:
+        return "<{}: with {} storage items and max_storage_size={}, base_state_lifetime={}>".format(
+            self.__class__.__name__,
+            len(self.storage),
+            self.max_storage_size,
+            self.base_state_lifetime,
+        )
+
+    def add_hasher[Event: BaseCute, HasherData](
+        self,
+        hasher: Hasher[Event, HasherData],
+        view: View,
+        /,
+    ) -> None:
+        self.storage[hasher] = LimitedDict(maxlimit=self.max_storage_size)
+        view.middlewares.insert(0, WaiterMiddleware(self, hasher))
+
+    async def drop_all(self) -> None:
+        for hasher in self.storage.copy():
+            for ident, short_state in self.storage[hasher].items():
+                if short_state.context:
+                    await self.drop(hasher, data=ident)
+                else:
+                    short_state.cancel_drop()
+                    await short_state.cancel()
+
+            del self.storage[hasher]
+
+    async def drop[Event: BaseCute, HasherData](
+        self,
+        hasher: Hasher[Event, HasherData],
+        data: HasherData,
+        *,
+        expired: bool = False,
+        **context: typing.Any,
+    ) -> None:
+        if hasher not in self.storage:
+            raise LookupError("No record of hasher {!r} found.".format(hasher))
+
+        waiter_id: typing.Hashable = hasher.get_hash_from_data(data).expect(
+            RuntimeError("Couldn't create hash from data"),
+        )
+        short_state = self.storage[hasher].pop(waiter_id, None)
+        if short_state is None:
+            raise LookupError("Waiter with identificator {} is not found for hasher {!r}.".format(waiter_id, hasher))
+
+        try:
+            if not expired:
+                short_state.cancel_drop()
+
+            context["short_state"] = short_state
+
+            if on_drop := short_state.actions.get("on_drop"):
+                await maybe_awaitable(bundle(on_drop, context)())
+        finally:
+            await short_state.cancel()
+
+    async def drop_state[Event: BaseCute, HasherData](
+        self,
+        short_state: ShortState[Event],
+        hasher: Hasher[Event, HasherData],
+        data: HasherData,
+        *,
+        expired: bool = False,
+        **context: typing.Any,
+    ) -> None:
+        preset_context = (
+            short_state.context.context.copy().as_dict() | context if short_state.context is not None else context
+        )
+
+        try:
+            await self.drop(hasher, data, expired=expired, **preset_context)
+        except Exception as e:
+            await logger.aerror("Error dropping state for hasher {!r}: {}", hasher, e)
+
+    async def drop_state_many[Event: BaseCute, ViewType: View, HasherData](
+        self,
+        short_state: ShortState[Event],
+        *hashers: HasherWithData[Event, ViewType, HasherData],
+        expired: bool = False,
+        **context: typing.Any,
+    ) -> None:
+        for hasher, _, data in hashers:
+            await self.drop_state(short_state, hasher, data, expired=expired, **context)
+
+    async def wait[Event: BaseCute, ViewType: View, HasherData](
+        self,
+        hasher: Hasher[Event, HasherData] | HasherWithData[Event, ViewType, HasherData],
+        view: ViewType | None = None,
+        data: HasherData = _NODATA,
+        *,
+        filter: ABCRule | None = None,
+        release: ABCRule | None = None,
+        lifetime: datetime.timedelta | float | None = None,
+        lifespan: Lifespan | None = None,
+        **actions: typing.Unpack[WaiterActions[Event]],
+    ) -> ShortStateContext[Event]:
+        if isinstance(lifetime, int | float):
+            lifetime = datetime.timedelta(seconds=lifetime)
+        elif lifetime is None:
+            lifetime = self.base_state_lifetime
+
+        lifespan = lifespan or Lifespan()
+        event = asyncio.Event()
+        short_state = ShortState(
+            event,
+            actions,
+            release=release,
+            filter=filter,
+            expiration=lifetime,
+        )
+
+        hasher, view, data = hasher if not isinstance(hasher, Hasher) else (hasher, view, data)
+
+        if data is _NODATA:
+            raise ValueError("Hasher requires data.")
+
+        if view is None:
+            raise ValueError("Hasher requires view.")
+
+        waiter_hash = hasher.get_hash_from_data(data).expect(RuntimeError("Hasher couldn't create hash."))
+
+        if hasher not in self.storage:
+            self.add_hasher(hasher, view)
+
+        if (deleted_short_state := self.storage[hasher].set(waiter_hash, short_state)) is not None:
+            deleted_short_state.cancel_drop()
+            await deleted_short_state.cancel()
+
+        async with lifespan:
+            short_state.schedule_drop(self.drop_state, hasher, data, lifetime=lifetime)
+            await event.wait()
+            short_state.cancel_drop()
+
+        self.storage[hasher].pop(waiter_hash, None)
+
+        if short_state.context is None:
+            raise LookupError("No context in short_state.")
+        return short_state.context
+
+    async def wait_many[Event: BaseCute[typing.Any], ViewType: View, Data, *Ts](
+        self,
+        *hashers: HasherWithData[Event, ViewType, Data],
+        filter: ABCRule | None = None,
+        release: ABCRule | None = None,
+        lifetime: datetime.timedelta | float | None = None,
+        lifespan: Lifespan | None = None,
+        unpack: ContextUnpackProto[*Ts] = unpack_to_context,
+        **actions: typing.Unpack[WaiterActions[Event]],
+    ) -> tuple[HasherWithData[Event, ViewType, Data], Event, *Ts]:
+        if isinstance(lifetime, int | float):
+            lifetime = datetime.timedelta(seconds=lifetime)
+        elif lifetime is None:
+            lifetime = self.base_state_lifetime
+
+        lifespan = lifespan or Lifespan()
+        event = asyncio.Event()
+        short_state = ShortState(
+            event,
+            actions,
+            release=release,
+            filter=filter,
+            expiration=lifetime,
+        )
+        waiter_hashes: dict[Hasher[Event, Data], typing.Hashable] = {}
+
+        for hasher, view, data in hashers:
+            waiter_hash = hasher.get_hash_from_data(data).expect(RuntimeError("Hasher couldn't create hash."))
+
+            if hasher not in self.storage:
+                self.add_hasher(hasher, view)
+
+            if (deleted_short_state := self.storage[hasher].set(waiter_hash, short_state)) is not None:
+                deleted_short_state.cancel_drop()
+                await deleted_short_state.cancel()
+
+            waiter_hashes[hasher] = waiter_hash
+
+        async with lifespan:
+            short_state.schedule_drop_many(self.drop_state_many, *hashers, lifetime=lifetime)
+            await event.wait()
+            short_state.cancel_drop()
+
+        if short_state.context is None:
+            raise LookupError("No context in short_state.")
+
+        initiator = short_state.context.context.get("initiator")
+        if initiator is None:
+            raise LookupError("Initiator not found in short_state context.")
+
+        for hasher, waiter_hash in waiter_hashes.items():
+            self.storage[hasher].pop(waiter_hash, None)
+
+        return (
+            initiator,
+            short_state.context.event,
+            *unpack(short_state.context.context),
+        )
+
+
+__all__ = ("WaiterMachine",)

telegrinder/bot/dispatch/waiter_machine/middleware.py

@@ -0,0 +1,77 @@
+import datetime
+import typing
+
+from telegrinder.api.api import API
+from telegrinder.bot.cute_types.update import UpdateCute
+from telegrinder.bot.dispatch.context import Context
+from telegrinder.bot.dispatch.handler.func import FuncHandler
+from telegrinder.bot.dispatch.middleware.abc import ABCMiddleware
+from telegrinder.bot.dispatch.process import check_rule
+from telegrinder.bot.dispatch.waiter_machine.short_state import ShortStateContext
+from telegrinder.modules import logger
+from telegrinder.types.objects import Update
+
+if typing.TYPE_CHECKING:
+    from telegrinder.bot.dispatch.waiter_machine.hasher import Hasher
+    from telegrinder.bot.dispatch.waiter_machine.machine import WaiterMachine
+    from telegrinder.bot.dispatch.waiter_machine.short_state import ShortState
+
+
+class WaiterMiddleware(ABCMiddleware):
+    def __init__(
+        self,
+        machine: WaiterMachine,
+        hasher: Hasher[typing.Any, typing.Any],
+    ) -> None:
+        self.machine = machine
+        self.hasher = hasher
+
+    async def pre(self, update_cute: UpdateCute, raw_update: Update, api: API, ctx: Context) -> bool:
+        if self.hasher not in self.machine.storage:
+            return True
+
+        event = update_cute.incoming_update
+        key = self.hasher.get_hash_from_data_from_event(event)
+        if not key:
+            await logger.ainfo("Unable to get hash from event with hasher {!r}", self.hasher)
+            return True
+
+        short_state: ShortState | None = self.machine.storage[self.hasher].get(key.unwrap())
+        if not short_state:
+            return True
+
+        # Just ignore update if state expired, so it will be dropped automatically by WaiterMachine
+        if short_state.expiration_date is not None and datetime.datetime.now() >= short_state.expiration_date:
+            return True
+
+        preset_context = Context(short_state=short_state)
+        if short_state.context is not None:
+            preset_context |= short_state.context.context
+
+        if short_state.filter is not None and not await check_rule(short_state.filter, preset_context):
+            await logger.adebug("Filter rule {!r} failed!", short_state.filter)
+            return True
+
+        handler = FuncHandler(
+            function=self.pass_runtime,
+            rules=(short_state.release,) if short_state.release is not None else None,
+            preset_context=preset_context,
+        )
+
+        if not await handler.run(api, raw_update, ctx) and (on_miss := short_state.actions.get("on_miss")) is not None:
+            await on_miss.run(api, raw_update, ctx)
+
+        return False
+
+    async def pass_runtime(
+        self,
+        event: UpdateCute,
+        ctx: Context,
+        short_state: ShortState,
+    ) -> None:
+        ctx.initiator = self.hasher
+        short_state.context = ShortStateContext(event.incoming_update, ctx)
+        short_state.event.set()
+
+
+__all__ = ("WaiterMiddleware",)

telegrinder/bot/dispatch/waiter_machine/short_state.py

@@ -0,0 +1,105 @@
+import asyncio
+import dataclasses
+import datetime
+import typing
+from functools import partial
+
+from telegrinder.bot.dispatch.context import Context
+from telegrinder.bot.rules.abc import ABCRule
+from telegrinder.tools.aio import cancel_future
+from telegrinder.tools.global_context.builtin_context import TelegrinderContext
+
+if typing.TYPE_CHECKING:
+    from telegrinder.bot.cute_types.base import BaseCute
+    from telegrinder.bot.dispatch.view.base import View
+    from telegrinder.bot.dispatch.waiter_machine.actions import WaiterActions
+    from telegrinder.bot.dispatch.waiter_machine.hasher import Hasher
+    from telegrinder.bot.dispatch.waiter_machine.machine import HasherWithData
+
+TELEGRINDER_CONTEXT: typing.Final = TelegrinderContext()
+
+
+def _wrap_async_fn[**P](
+    fn: typing.Callable[P, typing.Any],
+    /,
+) -> typing.Callable[P, None]:
+    def wrapper(*args: P.args, **kwargs: P.kwargs) -> None:
+        TELEGRINDER_CONTEXT.loop_wrapper.add_task(fn(*args, **kwargs))
+
+    return wrapper
+
+
+class ShortStateContext[Event: BaseCute[typing.Any] = typing.Any](typing.NamedTuple):
+    event: Event
+    context: Context
+
+
+@dataclasses.dataclass
+class ShortState[Event: BaseCute[typing.Any] = typing.Any]:
+    event: asyncio.Event
+    actions: WaiterActions[Event]
+
+    release: ABCRule | None = dataclasses.field(
+        default=None,
+        kw_only=True,
+    )
+    filter: ABCRule | None = dataclasses.field(
+        default=None,
+        kw_only=True,
+    )
+    expiration: dataclasses.InitVar[datetime.timedelta] = dataclasses.field(
+        kw_only=True,
+    )
+
+    expiration_date: datetime.datetime | None = dataclasses.field(init=False, kw_only=True)
+    creation_date: datetime.datetime = dataclasses.field(init=False)
+    context: ShortStateContext[Event] | None = dataclasses.field(default=None, init=False, kw_only=True)
+    drop_timer: asyncio.TimerHandle | None = dataclasses.field(default=None, init=False)
+
+    def __post_init__(self, expiration: datetime.timedelta) -> None:
+        self.lifetime = expiration
+        self.creation_date = datetime.datetime.now()
+        self.expiration_date = (self.creation_date + expiration) if expiration is not None else None
+
+    def schedule_drop[HasherData](
+        self,
+        dropper: typing.Callable[..., typing.Any],
+        /,
+        hasher: Hasher[Event, HasherData],
+        data: HasherData,
+        **context: typing.Any,
+    ) -> None:
+        self.drop_timer = asyncio.get_running_loop().call_later(
+            delay=self.lifetime.total_seconds(),
+            callback=partial(_wrap_async_fn(dropper), self, hasher, data, expired=True, **context),
+        )
+
+    def schedule_drop_many[ViewType: View, HasherData](
+        self,
+        dropper: typing.Callable[..., typing.Any],
+        /,
+        *hashers: HasherWithData[Event, ViewType, HasherData],
+        **context: typing.Any,
+    ) -> None:
+        self.drop_timer = asyncio.get_running_loop().call_later(
+            delay=self.lifetime.total_seconds(),
+            callback=partial(_wrap_async_fn(dropper), self, *hashers, expired=True, **context),
+        )
+
+    def cancel_drop(self) -> None:
+        if self.drop_timer is not None:
+            self.drop_timer.cancel()
+            self.drop_timer = None
+
+    async def cancel(self) -> None:
+        waiters = typing.cast(
+            "typing.Iterable[asyncio.Future[typing.Any]]",
+            self.event._waiters,  # type: ignore
+        )
+
+        for future in waiters:
+            if not future.cancelled():
+                await cancel_future(future)
+
+
+__all__ = ("ShortState", "ShortStateContext")

telegrinder/bot/polling/__init__.py

@@ -0,0 +1,4 @@
+from .abc import ABCPolling
+from .polling import Polling
+
+__all__ = ("ABCPolling", "Polling")

telegrinder/bot/polling/abc.py

@@ -0,0 +1,25 @@
+import typing
+from abc import ABC, abstractmethod
+
+import msgspec
+
+from telegrinder.types.objects import Update
+
+
+class ABCPolling(ABC):
+    offset: int
+
+    @abstractmethod
+    async def get_updates(self) -> list[msgspec.Raw]:
+        pass
+
+    @abstractmethod
+    async def listen(self) -> typing.AsyncGenerator[list[Update], None]:
+        yield []
+
+    @abstractmethod
+    def stop(self) -> None:
+        pass
+
+
+__all__ = ("ABCPolling",)

telegrinder/bot/polling/error_handler.py

@@ -0,0 +1,93 @@
+import asyncio
+import sys
+import typing
+
+from telegrinder.api.error import APIServerError, InvalidTokenError
+from telegrinder.modules import logger
+from telegrinder.tools.aio import maybe_awaitable
+
+if typing.TYPE_CHECKING:
+    from telegrinder.bot.polling.polling import Polling
+
+
+class ErrorHandler:
+    _handlers: dict[type[BaseException], typing.Callable[[BaseException], typing.Any]]
+
+    __slots__ = ("_polling", "_handlers")
+
+    def __init__(self, polling: Polling, /) -> None:
+        self._polling = polling
+        self._handlers = {  # type: ignore
+            InvalidTokenError: self._handle_invalid_token_error,
+            asyncio.CancelledError: self._handle_cancelled_error,
+            APIServerError: self._handle_api_server_error,
+            **{e: self._handle_connection_timeout_error for e in polling.api.http.CONNECTION_TIMEOUT_ERRORS},
+            **{e: self._handle_client_connection_error for e in polling.api.http.CLIENT_CONNECTION_ERRORS},
+        }
+
+    async def handle(self, error: BaseException) -> bool:
+        error_class = type(error)
+
+        if error_class is SystemExit:
+            self._handle_system_exit(error)  # type: ignore
+
+        if error_class not in self._handlers:
+            return False
+
+        try:
+            await maybe_awaitable(self._handlers[error_class](error))
+            return True
+        except SystemExit as sys_exit_err:
+            await self._handle_system_exit(sys_exit_err)
+
+    async def _handle_system_exit(self, error: SystemExit) -> typing.NoReturn:
+        await logger.aerror("Forced exit from the program with code {}.", error.code)
+        raise error from None
+
+    async def _handle_invalid_token_error(
+        self,
+        error: InvalidTokenError,
+    ) -> typing.NoReturn:
+        await logger.aerror("{}", error)
+        self._polling.stop()
+        sys.exit(3)
+
+    async def _handle_cancelled_error(self, _: asyncio.CancelledError) -> None:
+        await logger.ainfo("Caught cancel, stopping polling")
+        self._polling.stop()
+
+    async def _handle_connection_timeout_error(self, _: BaseException) -> None:
+        if self._polling.reconnects_counter > self._polling.max_reconnects:
+            await logger.aerror(
+                "Failed to reconnect to Telegram API server after {} attempts, stopping polling",
+                self._polling.max_reconnects,
+            )
+            self._polling.stop()
+            sys.exit(6)
+
+        await logger.awarning(
+            "Server disconnected, waiting {} seconds to reconnect...",
+            self._polling.reconnect_after,
+        )
+        await asyncio.sleep(self._polling.reconnect_after)
+
+    async def _handle_client_connection_error(self, _: BaseException) -> None:
+        await logger.aerror(
+            "Client connection failed, attempt to reconnect after {} seconds...",
+            self._polling.reconnect_after,
+        )
+        await asyncio.sleep(self._polling.reconnect_after)
+
+    async def _handle_api_server_error(
+        self,
+        error: APIServerError,
+    ) -> None:
+        if error.retry_after is None:
+            await logger.aerror("{}", error)
+            sys.exit(9)
+
+        await logger.aerror("{}, waiting {} seconds to the next request...", error, error.retry_after)
+        await asyncio.sleep(error.retry_after)
+
+
+__all__ = ("ErrorHandler",)