diff options
Diffstat (limited to 'libs')
| -rw-r--r-- | libs/__init__.py | 7 | ||||
| -rw-r--r-- | libs/fsm.py | 23 | ||||
| -rw-r--r-- | libs/msg.py | 25 | ||||
| -rw-r--r-- | libs/storage.py | 75 |
4 files changed, 130 insertions, 0 deletions
diff --git a/libs/__init__.py b/libs/__init__.py new file mode 100644 index 0000000..d8ed122 --- /dev/null +++ b/libs/__init__.py | |||
| @@ -0,0 +1,7 @@ | |||
| 1 | from . import fsm, msg, storage | ||
| 2 | |||
| 3 | __all__ = [ | ||
| 4 | "storage", | ||
| 5 | "fsm", | ||
| 6 | "msg", | ||
| 7 | ] | ||
diff --git a/libs/fsm.py b/libs/fsm.py new file mode 100644 index 0000000..00ccacb --- /dev/null +++ b/libs/fsm.py | |||
| @@ -0,0 +1,23 @@ | |||
| 1 | from contextlib import asynccontextmanager | ||
| 2 | |||
| 3 | from aiogram.fsm.context import FSMContext | ||
| 4 | from pydantic.main import BaseModel | ||
| 5 | from typing_extensions import AsyncGenerator | ||
| 6 | |||
| 7 | |||
| 8 | async def set_data(state: FSMContext, model: BaseModel) -> None: | ||
| 9 | await state.set_data(model.model_dump()) | ||
| 10 | |||
| 11 | |||
| 12 | async def get_data[T: BaseModel](state: FSMContext, model_type: type[T]) -> T: | ||
| 13 | return model_type.model_validate(await state.get_data()) | ||
| 14 | |||
| 15 | |||
| 16 | @asynccontextmanager | ||
| 17 | async def edit_data[T: BaseModel]( | ||
| 18 | state: FSMContext, | ||
| 19 | model_type: type[T], | ||
| 20 | ) -> AsyncGenerator[T]: | ||
| 21 | model = await get_data(state, model_type) | ||
| 22 | yield model | ||
| 23 | await set_data(state, model) | ||
diff --git a/libs/msg.py b/libs/msg.py new file mode 100644 index 0000000..9bcc52a --- /dev/null +++ b/libs/msg.py | |||
| @@ -0,0 +1,25 @@ | |||
| 1 | import asyncio | ||
| 2 | from typing import AsyncGenerator, Iterable | ||
| 3 | |||
| 4 | from aiogram import Bot | ||
| 5 | from aiogram.exceptions import TelegramAPIError, TelegramRetryAfter | ||
| 6 | |||
| 7 | from models import RichText | ||
| 8 | |||
| 9 | |||
| 10 | async def publish( | ||
| 11 | bot: Bot, | ||
| 12 | users: Iterable[int], | ||
| 13 | rich_text: RichText, | ||
| 14 | ) -> AsyncGenerator[int]: | ||
| 15 | for n, user_id in enumerate(users, start=1): | ||
| 16 | for _ in range(5): | ||
| 17 | try: | ||
| 18 | await rich_text.send(bot, user_id) | ||
| 19 | break | ||
| 20 | except TelegramRetryAfter as e: | ||
| 21 | await asyncio.sleep(e.retry_after + 1) | ||
| 22 | except TelegramAPIError: | ||
| 23 | await asyncio.sleep(5) | ||
| 24 | |||
| 25 | yield n | ||
diff --git a/libs/storage.py b/libs/storage.py new file mode 100644 index 0000000..6220cfc --- /dev/null +++ b/libs/storage.py | |||
| @@ -0,0 +1,75 @@ | |||
| 1 | from pathlib import Path | ||
| 2 | from typing import Any, Mapping | ||
| 3 | |||
| 4 | from aiofiles import open as open | ||
| 5 | from aiogram.fsm.state import State | ||
| 6 | from aiogram.fsm.storage.base import ( | ||
| 7 | BaseStorage, | ||
| 8 | DefaultKeyBuilder, | ||
| 9 | KeyBuilder, | ||
| 10 | StateType, | ||
| 11 | StorageKey, | ||
| 12 | ) | ||
| 13 | from pydantic import TypeAdapter | ||
| 14 | from pydantic.main import BaseModel | ||
| 15 | |||
| 16 | |||
| 17 | class Record(BaseModel): | ||
| 18 | data: dict[str, Any] = {} | ||
| 19 | state: str | None = None | ||
| 20 | |||
| 21 | |||
| 22 | class JsonStorage(BaseStorage): | ||
| 23 | file_path: Path | ||
| 24 | records: dict[str, Record] | ||
| 25 | records_adapter: TypeAdapter | ||
| 26 | key_builder: KeyBuilder | ||
| 27 | |||
| 28 | def __init__(self, file_path: Path, key_builder: KeyBuilder | None = None) -> None: | ||
| 29 | self.file_path = file_path | ||
| 30 | self.records = {} | ||
| 31 | self.records_adapter = TypeAdapter(dict[str, Record]) | ||
| 32 | self.key_builder = DefaultKeyBuilder() if key_builder is None else key_builder | ||
| 33 | |||
| 34 | async def read(self) -> None: | ||
| 35 | async with open(self.file_path, "rb") as file: | ||
| 36 | json = await file.read() | ||
| 37 | self.records = self.records_adapter.validate_json(json) | ||
| 38 | |||
| 39 | async def flush(self) -> None: | ||
| 40 | async with open(self.file_path, "wb") as file: | ||
| 41 | json = self.records_adapter.dump_json(self.records) | ||
| 42 | await file.write(json) | ||
| 43 | |||
| 44 | async def get_record(self, key: StorageKey) -> Record: | ||
| 45 | await self.read() | ||
| 46 | record_key = self.key_builder.build(key) | ||
| 47 | if record_key not in self.records: | ||
| 48 | self.records[record_key] = Record() | ||
| 49 | return self.records[record_key] | ||
| 50 | |||
| 51 | async def set_state(self, key: StorageKey, state: StateType = None) -> None: | ||
| 52 | record = await self.get_record(key) | ||
| 53 | record.state = state.state if isinstance(state, State) else state | ||
| 54 | await self.flush() | ||
| 55 | |||
| 56 | async def get_state(self, key: StorageKey) -> str | None: | ||
| 57 | record = await self.get_record(key) | ||
| 58 | return record.state | ||
| 59 | |||
| 60 | async def set_data(self, key: StorageKey, data: Mapping[str, Any]) -> None: | ||
| 61 | if not isinstance(data, dict): | ||
| 62 | raise TypeError( | ||
| 63 | f"Data must be a dict or dict-like object, got {type(data).__name__}", | ||
| 64 | data, | ||
| 65 | ) | ||
| 66 | record = await self.get_record(key) | ||
| 67 | record.data = data.copy() | ||
| 68 | await self.flush() | ||
| 69 | |||
| 70 | async def get_data(self, key: StorageKey) -> dict[str, Any]: | ||
| 71 | record = await self.get_record(key) | ||
| 72 | return record.data | ||
| 73 | |||
| 74 | async def close(self) -> None: | ||
| 75 | await self.flush() | ||
