aboutsummaryrefslogtreecommitdiff
path: root/libs/fsm.py
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)