refactor: dialog_extensions and handles

Author: upconett

src/bot/dialog_extension/__init__.py

@@ -0,0 +1,3 @@
+from src.bot.dialog_extension.extended_dialog_manager import ExtendedDialogManager, extend_dialog
+
+__all__ = ["extend_dialog", "ExtendedDialogManager"]

src/bot/dialog_extension/dialog_wrapper.py

@@ -0,0 +1,106 @@
+from typing import Any
+
+from aiogram.fsm.state import State
+from aiogram_dialog import DialogManager, ShowMode
+from aiogram_dialog.api.entities import AccessSettings, ChatEvent, Context, Stack, StartMode
+from aiogram_dialog.api.protocols.manager import BaseDialogManager, UnsetId
+
+
+class DialogManagerWrapper(DialogManager):
+    original: DialogManager
+
+    def __init__(self, dialog_manager: DialogManager):
+        self.original = dialog_manager
+
+    async def done(self, result: Any = None, show_mode: ShowMode | None = None) -> None:
+        return await self.original.done(result, show_mode)
+
+    async def start(
+        self,
+        state: State,
+        data: dict | list | int | str | float | None = None,
+        mode: StartMode = StartMode.NORMAL,
+        show_mode: ShowMode | None = None,
+        access_settings: AccessSettings | None = None,
+    ) -> None:
+        return await self.original.start(state, data, mode, show_mode, access_settings)
+
+    async def switch_to(self, state: State, show_mode: ShowMode | None = None) -> None:
+        return await self.original.switch_to(state, show_mode)
+
+    async def update(self, data: dict, show_mode: ShowMode | None = None) -> None:
+        return await self.original.update(data, show_mode)
+
+    def bg(
+        self,
+        user_id: int | None = None,
+        chat_id: int | None = None,
+        stack_id: str | None = None,
+        thread_id: int | None | UnsetId = UnsetId.UNSET,
+        business_connection_id: str | None | UnsetId = UnsetId.UNSET,
+        load: bool = False,
+    ) -> BaseDialogManager:
+        return self.original.bg(user_id, chat_id, stack_id, thread_id, business_connection_id, load)
+
+    @property
+    def event(self) -> ChatEvent:
+        return self.original.event
+
+    async def mark_closed(self) -> None:
+        return await self.original.mark_closed()
+
+    @property
+    def middleware_data(self) -> dict:
+        return self.original.middleware_data
+
+    @property
+    def dialog_data(self) -> dict:
+        return self.original.dialog_data
+
+    @property
+    def start_data(self) -> dict | list | int | str | float | None:
+        return self.original.start_data
+
+    @property
+    def show_mode(self) -> ShowMode:
+        return self.original.show_mode
+
+    @show_mode.setter
+    def show_mode(self, show_mode: ShowMode) -> None:
+        self.original.show_mode = show_mode
+
+    def is_preview(self) -> bool:
+        return self.original.is_preview()
+
+    async def show(self, show_mode: ShowMode | None = None) -> None:
+        return await self.original.show(show_mode)
+
+    async def answer_callback(self) -> None:
+        return await self.original.answer_callback()
+
+    def current_context(self) -> Context:
+        return self.original.current_context()
+
+    def has_context(self) -> bool:
+        return self.original.has_context()
+
+    def current_stack(self) -> Stack:
+        return self.original.current_stack()
+
+    async def next(self, show_mode: ShowMode | None = None) -> None:
+        return await self.original.next(show_mode)
+
+    async def back(self, show_mode: ShowMode | None = None) -> None:
+        return await self.original.back(show_mode)
+
+    def find(self, widget_id) -> Any | None:
+        return self.original.find(widget_id)
+
+    async def reset_stack(self, remove_keyboard: bool = True) -> None:
+        return await self.original.reset_stack(remove_keyboard)
+
+    async def load_data(self) -> dict:
+        return await self.original.load_data()
+
+    async def close_manager(self) -> None:
+        return await self.original.close_manager()

src/bot/dialog_extension/extended_dialog_manager.py

@@ -0,0 +1,54 @@
+from aiogram import Bot
+from aiogram.types import Chat, Message
+from aiogram_dialog import DialogManager, ShowMode
+
+from src.bot.dto import *
+
+from .dialog_wrapper import DialogManagerWrapper
+from .extended_fsm_context import ExtendedFSMContext, extend_fsm_context
+
+
+class ExtendedDialogManager(DialogManagerWrapper):
+    @property
+    def state(self) -> ExtendedFSMContext:
+        return extend_fsm_context(self.middleware_data["state"])
+
+    @property
+    def bot(self) -> Bot:
+        return self.middleware_data["bot"]
+
+    @property
+    def chat(self) -> Chat:
+        return self.middleware_data["event_chat"]
+
+    async def track_message(self, message: Message):
+        """Track messages to delete later ("to_delete_list" in state)"""
+        await self.state.add_to_delete(message)
+
+    async def clear_messages(self):
+        """Delete tracked messages ("to_delete_id" in state)"""
+        to_delete_list = await self.state.get_to_delete_list()
+        print("TRYING TO CLEAR:", to_delete_list)
+        try:
+            await self.bot.delete_messages(self.chat.id, to_delete_list)
+        except Exception as e:
+            return print(f"clear_messages could not delete {len(to_delete_list)} messages, {e}")
+        await self.state.set_to_delete_list([])
+
+    async def switch_to_current(self, show_mode: ShowMode):
+        return await self.switch_to(self.current_context().state, show_mode=ShowMode.DELETE_AND_SEND)
+
+    async def answer_and_track(self, text: str, **kwargs):
+        to_delete = await self.bot.send_message(chat_id=self.chat.id, text=text, **kwargs)
+        await self.track_message(to_delete)
+
+    async def answer_and_retry(self, text: str, **kwargs):
+        to_delete = await self.bot.send_message(chat_id=self.chat.id, text=text, **kwargs)
+        await self.track_message(to_delete)
+        await self.switch_to_current(ShowMode.DELETE_AND_SEND)
+
+
+def extend_dialog(dialog_manager: DialogManager) -> ExtendedDialogManager:
+    if not isinstance(dialog_manager, ExtendedDialogManager):
+        dialog_manager = ExtendedDialogManager(dialog_manager)
+    return dialog_manager

src/bot/extended_fsm_context.py …dialog_extension/extended_fsm_context.pysrc/bot/extended_fsm_context.py renamed to src/bot/dialog_extension/extended_fsm_context.py src/bot/extended_fsm_context.py renamed to src/bot/dialog_extension/extended_fsm_context.py

@@ -56,7 +56,7 @@ async def add_to_delete(self, message: Message):
         await self.set_to_delete_list(to_delete_list)
 
 
-def extend(fsm_context: FSMContext) -> ExtendedFSMContext:
+def extend_fsm_context(fsm_context: FSMContext) -> ExtendedFSMContext:
     if not isinstance(fsm_context, ExtendedFSMContext):
         fsm_context = ExtendedFSMContext.wrap(fsm_context)
     return fsm_context

src/bot/dialogs/attendance/handles.py

@@ -2,8 +2,8 @@
 from aiogram_dialog import DialogManager, ShowMode
 from aiogram_dialog.widgets.input import MessageInput
 
+from src.bot.dialog_extension import extend_dialog
 from src.bot.dto import *
-from src.bot.extended_dialog_manager import extend
 from src.bot.user_errors import *
 from src.bot.utils import *
 from src.db.repositories import meetings_repo
@@ -14,8 +14,8 @@
 
 
 async def get_attendance_file_close(message: Message, _: MessageInput, manager: DialogManager):
-    manager = extend(manager)
-    await clear_messages(manager)
+    manager = extend_dialog(manager)
+    await manager.clear_messages()
     await message.delete()
     try:
         contents = await get_document_contents(message, manager)
@@ -33,8 +33,8 @@ async def get_attendance_file_close(message: Message, _: MessageInput, manager:
 
 
 async def get_attendance_file_resend(message: Message, _: MessageInput, manager: DialogManager):
-    manager = extend(manager)
-    await clear_messages(manager)
+    manager = extend_dialog(manager)
+    await manager.clear_messages()
     await message.delete()
     try:
         contents = await get_document_contents(message, manager)
@@ -52,7 +52,7 @@ async def get_attendance_file_resend(message: Message, _: MessageInput, manager:
 
 
 async def on_download_attendance(query: CallbackQuery, button: Button, manager: DialogManager):
-    manager = extend(manager)
+    manager = extend_dialog(manager)
     try:
         input_file = await get_attendance_file_to_download(manager)
         await query.answer("Sending attendance file...")
@@ -65,15 +65,14 @@ async def on_download_attendance(query: CallbackQuery, button: Button, manager:
 
 
 async def get_email_to_add(message: Message, _: MessageInput, manager: DialogManager):
-    manager = extend(manager)
-    await clear_messages(manager)
+    manager = extend_dialog(manager)
+    await manager.clear_messages()
     await message.delete()
     try:
         email = await extract_email(message)
+        await add_email_to_attendance(email, manager)
     except NoMessageText:
         return await manager.answer_and_retry("There's no text in your message, enter email of a person to add")
-    try:
-        await add_email_to_attendance(email, manager)
     except NoMeetingAttendance:
         return await manager.answer_and_retry("Somehow there is no attendance for this meeting, resend the file maybe")
     except EmailAlreadyPresent:

src/bot/dialogs/attendance/logic.py

@@ -1,7 +1,7 @@
-from aiogram.types import BufferedInputFile, InputFile
+from aiogram.types import BufferedInputFile, InputFile, Message
 from aiogram_dialog import DialogManager
 
-from src.bot.extended_dialog_manager import extend
+from src.bot.dialog_extension import extend_dialog
 from src.bot.scheduling import *
 from src.bot.user_errors import *
 from src.bot.utils import *
@@ -11,7 +11,7 @@
 
 
 async def get_document_contents(message: Message, manager: DialogManager) -> str:
-    manager = extend(manager)
+    manager = extend_dialog(manager)
     if not message.document:
         raise NoDocumentError()
     file = await manager.bot.get_file(message.document.file_id)
@@ -28,7 +28,7 @@ async def get_document_contents(message: Message, manager: DialogManager) -> str
 
 
 async def get_attendance_file_to_download(manager: DialogManager) -> InputFile:
-    manager = extend(manager)
+    manager = extend_dialog(manager)
     meeting = await manager.state.get_meeting()
     if not meeting.attendance:
         raise NoMeetingAttendance()
@@ -47,7 +47,7 @@ async def extract_email(message: Message) -> Email:
 
 
 async def add_email_to_attendance(email: Email, manager: DialogManager):
-    manager = extend(manager)
+    manager = extend_dialog(manager)
     async with manager.state.sync_meeting() as meeting:
         if not meeting.attendance:
             raise NoMeetingAttendance()

src/bot/dialogs/attendance/windows.py

@@ -32,7 +32,7 @@
 resend_ww = Window(
     Format('Resend attendance file for "{title}"'),
     Row(
-        SwitchTo(Const("Back"), state=AttendanceStates.init, id="to_attendance_init"),
+        SwitchTo(Const("Back"), state=AttendanceStates.init, id="to_attendance_init", on_click=handle_clear),
         Button(Const(" "), id="blank", on_click=None),
     ),
     MessageInput(get_attendance_file_resend),
@@ -42,7 +42,7 @@
 
 close_ww = Window(
     Format('Send attendance file to close "{title}"'),
-    Row(Cancel(Const("Back")), Button(Const(" "), id="blank", on_click=None)),
+    Row(Cancel(Const("Back"), on_click=handle_clear), Button(Const(" "), id="blank", on_click=None)),
     MessageInput(get_attendance_file_close),
     getter=meeting_info_getter,
     state=AttendanceStates.close,
@@ -51,7 +51,7 @@
 add_email_ww = Window(
     Format("Enter email of a person to add 👤"),
     Row(
-        SwitchTo(Const("Back"), state=AttendanceStates.init, id="to_attendance_init"),
+        SwitchTo(Const("Back"), state=AttendanceStates.init, id="to_attendance_init", on_click=handle_clear),
         Button(Const(" "), id="blank", on_click=None),
     ),
     MessageInput(get_email_to_add),

src/bot/dialogs/change_meeting/dialog_buttons.py

@@ -9,7 +9,7 @@
 BTN_BLANK = Button(Const(" "), id="blank")
 BTN_ROW_BACK = Row(Back(), BTN_BLANK)
 
-BTN_INIT = SwitchTo(Const("Back"), id="to_init", state=ChangeStates.init)
+BTN_INIT = SwitchTo(Const("Back"), id="to_init", state=ChangeStates.init, on_click=handle_clear)
 
 
 TUTORS_ASSIGN_SCROLLING_GROUP = ScrollingGroup(

src/bot/dialogs/change_meeting/getters.py

@@ -1,5 +1,6 @@
 from aiogram_dialog import DialogManager
 
+from src.bot.dialog_extension import extend_dialog
 from src.bot.dialogs.meetings.getters import meeting_info_getter
 from src.bot.dto import *
 from src.bot.filters import *
@@ -8,14 +9,9 @@
 
 
 async def meeting_info_with_tutors_getter(dialog_manager: DialogManager, **kwargs):
-    state = get_state(dialog_manager)
-    meeting = dto_to_meeting(await state.get_value("meeting"))
-
-    if not meeting:
-        raise ValueError("No Meeting in meeting_info_getter")
-
+    manager = extend_dialog(dialog_manager)
+    meeting = await manager.state.get_meeting()
     tutors = await tutors_repo.list()
-
     return {
         "title": meeting.title,
         "description": meeting.description,