1 files changed, 23 insertions, 0 deletions

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 @@
+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)