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)