dj-embe/service/fileManager.py

import re
from collections.abc import AsyncIterator
from os import path, stat
from time import time

from discord import Interaction

from entity import Entry, File
from framework import Downloader, Youtube


class FileManager:
    def __init__(
        self, youtube: Youtube, downloader: Downloader, downloadDirectory: str
    ) -> None:
        self.youtube = youtube
        self.downloader = downloader
        self.downloadDirectory = downloadDirectory
        self.youtubeVideoRegex = re.compile(
            r"(https:\/\/)?(www|music)\.youtube\.com\/(watch\?v=|shorts\/)\w*"
        )
        self.progressLastUpdate = 0.0

    def getFileName(self, entryId: str) -> str:
        return path.join(self.downloadDirectory, entryId)

    async def downloadProgress(
        self,
        interaction: Interaction,
        entry: Entry,
        current_size: float,
        total_size: float,
    ):
        current_size = current_size / 1024 / 1024
        total_size = total_size / 1024 / 1024
        if time() - self.progressLastUpdate > 1 and total_size > 10:
            self.progressLastUpdate = time()
            await interaction.edit_original_response(
                content=f"Downloading {entry.title.name} [%.2f/%.2f Mo]"
                % (current_size, total_size)
            )

    async def download(
        self, interaction: Interaction, entries: list[Entry]
    ) -> AsyncIterator[File]:
        for entry in entries:
            print("DOWNLOAD", entry.source)
            entryId = self.youtube.getId(entry.source)

            fileName = self.getFileName(entryId)
            if not path.isfile(fileName) and isinstance(entry.source, str):
                if self.youtubeVideoRegex.match(entry.source) is not None:
                    yield await self.youtube.download(
                        entry.source,
                        fileName,
                        lambda current_size, total_size: self.downloadProgress(
                            interaction, entry, current_size, total_size
                        ),
                    )
                else:
                    yield await self.downloader.download(
                        entry.source,
                        fileName,
                        lambda current_size, total_size: self.downloadProgress(
                            interaction, entry, current_size, total_size
                        ),
                    )

            yield File(fileName, stat(fileName).st_size)