51 lines
2.7 KiB
Python
51 lines
2.7 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):
|
||
# latency - задержка в секундах для сбора частей альбома
|
||
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:
|
||
# Если у сообщения нет media_group_id, это не альбом -> пропускаем дальше как обычно
|
||
if not event.media_group_id:
|
||
return await handler(event, data)
|
||
|
||
group_id = event.media_group_id
|
||
|
||
try:
|
||
# Если этот альбом мы еще не видели (первое сообщение из пачки)
|
||
if group_id not in self.album_data:
|
||
self.album_data[group_id] = [event] # Создаем список
|
||
await asyncio.sleep(self.latency) # Ждем остальные части
|
||
|
||
# После ожидания кладем собранный список в data
|
||
# Теперь в хендлере будет доступен аргумент 'album'
|
||
data["album"] = self.album_data[group_id]
|
||
|
||
# Вызываем хендлер ОДИН раз
|
||
return await handler(event, data)
|
||
|
||
else:
|
||
# Если альбом уже собирается, просто добавляем сообщение в список
|
||
# и НЕ вызываем хендлер (прерываем цепочку для этого сообщения)
|
||
self.album_data[group_id].append(event)
|
||
return
|
||
|
||
finally:
|
||
# Чистим память после обработки, если это был "главный" поток обработки
|
||
if group_id in self.album_data and len(self.album_data[group_id]) > 1:
|
||
# Маленький хак: удаляем только если обработчик завершился
|
||
# Проверка len нужна, чтобы не удалить раньше времени в параллельных тасках,
|
||
# но корректнее просто удалять в блоке первого сообщения.
|
||
if event == self.album_data[group_id][0]:
|
||
del self.album_data[group_id] |