This commit is contained in:
xds
2026-02-03 09:14:45 +03:00
parent 4f533edf50
commit 3497f963fb
5 changed files with 122 additions and 81 deletions

View File

@@ -6,7 +6,6 @@ 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]] = {}
@@ -16,36 +15,33 @@ class AlbumMiddleware(BaseMiddleware):
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) # Ждем остальные части
# Если это первое сообщение группы
if group_id not in self.album_data:
self.album_data[group_id] = [event]
try:
# Ждем сбора остальных частей
await asyncio.sleep(self.latency)
# После ожидания кладем собранный список в data
# Теперь в хендлере будет доступен аргумент 'album'
data["album"] = self.album_data[group_id]
# Проверяем, что ключ все еще существует (на всякий случай)
if group_id in self.album_data:
# Передаем собранный альбом в хендлер
# Сортируем по message_id, чтобы порядок был верным
self.album_data[group_id].sort(key=lambda x: x.message_id)
data["album"] = self.album_data[group_id]
return await handler(event, data)
# Вызываем хендлер ОДИН раз
return await handler(event, data)
finally:
# ЧИСТКА: Удаляем всегда, если это "головной" поток, который создал запись
# Проверяем, что мы удаляем именно то, что создали, и ключ существует
if group_id in self.album_data and self.album_data[group_id][0] == event:
del self.album_data[group_id]
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]
else:
# Если группа уже собирается - просто добавляем и выходим
self.album_data[group_id].append(event)
return