70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
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)
|