47 lines
1.9 KiB
Python
47 lines
1.9 KiB
Python
import asyncio
|
||
from typing import Any, Dict, List, Union, Callable, Awaitable
|
||
from aiogram import BaseMiddleware
|
||
from aiogram.types import Message
|
||
|
||
|
||
class AlbumMiddleware(BaseMiddleware):
|
||
def __init__(self, latency: float = 0.5):
|
||
self.latency = latency
|
||
self.album_data: Dict[str, List[Message]] = {}
|
||
|
||
async def __call__(
|
||
self,
|
||
handler: Callable[[Message, Dict[str, Any]], Awaitable[Any]],
|
||
event: Message,
|
||
data: Dict[str, Any]
|
||
) -> Any:
|
||
if not event.media_group_id:
|
||
return await handler(event, data)
|
||
|
||
group_id = event.media_group_id
|
||
|
||
# Если это первое сообщение группы
|
||
if group_id not in self.album_data:
|
||
self.album_data[group_id] = [event]
|
||
try:
|
||
# Ждем сбора остальных частей
|
||
await asyncio.sleep(self.latency)
|
||
|
||
# Проверяем, что ключ все еще существует
|
||
if group_id in self.album_data:
|
||
# Передаем собранный альбом в хендлер
|
||
# Сортируем по message_id, чтобы порядок был верным
|
||
current_album = self.album_data[group_id]
|
||
current_album.sort(key=lambda x: x.message_id)
|
||
data["album"] = current_album
|
||
return await handler(event, data)
|
||
|
||
finally:
|
||
# ЧИСТКА: Удаляем запись после обработки или таймаута
|
||
# Используем pop() с дефолтом, чтобы избежать KeyError
|
||
self.album_data.pop(group_id, None)
|
||
|
||
else:
|
||
# Если группа уже собирается - просто добавляем и выходим
|
||
self.album_data[group_id].append(event)
|
||
return |