blob: 00ccacbd43cd448d855506467b5aa17a10c3fdba (
plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
|
from contextlib import asynccontextmanager
from aiogram.fsm.context import FSMContext
from pydantic.main import BaseModel
from typing_extensions import AsyncGenerator
async def set_data(state: FSMContext, model: BaseModel) -> None:
await state.set_data(model.model_dump())
async def get_data[T: BaseModel](state: FSMContext, model_type: type[T]) -> T:
return model_type.model_validate(await state.get_data())
@asynccontextmanager
async def edit_data[T: BaseModel](
state: FSMContext,
model_type: type[T],
) -> AsyncGenerator[T]:
model = await get_data(state, model_type)
yield model
await set_data(state, model)
|