Handle files in media group using aiogram

Question

I need a sustainable method to handle files wrapped in media group.

My function handle_files is waiting for media files. When user uploads media file it goes through the series of different checks. If it passes all tests (size restriction, formats restriction) the media file is downloaded and processed.

It looks like this:

async def handle_files(message: types.Message, state: FSMContext):
    user_data = await state.get_data()
    locale = user_data['locale']
    list_of_files = user_data['list_of_files']

    try:
        file = message.document
        file_name = file['file_name']
    except Exception as e:
        await message.answer('Error while downloading file')
        return None
    file_name = unidecode(file_name)
    file_size = file['file_size']
    if file_size >= 20971520:
        await message.answer('File is too big')
        return None
    invalid_format, formats = check_invalid_format(file_name, function)
    if invalid_format:
        await message.answer(file_name + 'has unsupported format. Supported formats: ' + ', '.join(formats))
        return None
    output_folder = os.path.join('temp', str(message.from_user.id))
    if not os.path.exists(output_folder):
        os.makedirs(output_folder, exist_ok=True)
    file_path = os.path.join(output_folder, file_name)

    await file.download(destination_file=file_path)

    list_of_files.append(file_path)
    await state.update_data(list_of_files=list_of_files)

    await message.answer('Added files: '.format('\n'.join(list_of_files)))

Working with separate files looks fine. After downloading files user gets a list of added files one by one.

But the problem is that when user uploads files in media group one file overrides another. It looks like that.

So, only one file is appending to the list list_of_files that prevents me from processing both files.

I tried to solve the problem by initializing user_data dictionary one more time:

    ...
    await file.download(destination_file=file_path)
    user_data = await state.get_data()
    list_of_files = user_data['list_of_files']
    list_of_files.append(file_path)
    await state.update_data(list_of_files=list_of_files)
    ...

It solved one part of my problem but this solution is not elegant and, supposedly, not quite sustainable. The message is duplicated.

I need that after uploading media group user gets one message containing the list of all files from this media group.

I suppose that the problem here is linked with asyncio. But I've spent already a lot of time but the solution haven't been found. Looking for you help.

Answer 1

check this link. Here you can find the middleware, which will return list of messages from media group, and then you could use for loop to handle them.

import asyncio
from typing import List, Union

from aiogram import Bot, Dispatcher, executor, types
from aiogram.dispatcher.handler import CancelHandler
from aiogram.dispatcher.middlewares import BaseMiddleware

bot = Bot(token="TOKEN_HERE")  # Place your token here
dp = Dispatcher(bot)


class AlbumMiddleware(BaseMiddleware):
    """This middleware is for capturing media groups."""

    album_data: dict = {}

    def __init__(self, latency: Union[int, float] = 0.01):
        """
        You can provide custom latency to make sure
        albums are handled properly in highload.
        """
        self.latency = latency
        super().__init__()

    async def on_process_message(self, message: types.Message, data: dict):
        if not message.media_group_id:
            return

        try:
            self.album_data[message.media_group_id].append(message)
            raise CancelHandler()  # Tell aiogram to cancel handler for this group element
        except KeyError:
            self.album_data[message.media_group_id] = [message]
            await asyncio.sleep(self.latency)

            message.conf["is_last"] = True
            data["album"] = self.album_data[message.media_group_id]

    async def on_post_process_message(self, message: types.Message, result: dict, data: dict):
        """Clean up after handling our album."""
        if message.media_group_id and message.conf.get("is_last"):
            del self.album_data[message.media_group_id]


@dp.message_handler(is_media_group=True, content_types=types.ContentType.ANY)
async def handle_albums(message: types.Message, album: List[types.Message]):
    """This handler will receive a complete album of any type."""
    media_group = types.MediaGroup()
    for obj in album:
        if obj.photo:
            file_id = obj.photo[-1].file_id
        else:
            file_id = obj[obj.content_type].file_id

        try:
            # We can also add a caption to each file by specifying `"caption": "text"`
            media_group.attach({"media": file_id, "type": obj.content_type})
        except ValueError:
            return await message.answer("This type of album is not supported by aiogram.")

    await message.answer_media_group(media_group)


if __name__ == "__main__":
    dp.middleware.setup(AlbumMiddleware())
    executor.start_polling(dp, skip_updates=True)