diff --git a/app/utils/youtube.py b/app/utils/youtube.py index 3d97236..a5a4199 100644 --- a/app/utils/youtube.py +++ b/app/utils/youtube.py @@ -1,6 +1,9 @@ -from typing import Optional +import subprocess +from typing import Literal from dataclasses import dataclass +from app.utils.uploader import ChunkUploadBackend + @dataclass class MediaContent: @@ -9,9 +12,10 @@ class MediaContent: chunk_size: int -class YtDlpInfo: - def __init__(self, url: str): +class YtDlpManager: + def __init__(self, url: str, backend: ChunkUploadBackend): self.url = url + self.backend = backend self.info = None self._extract_info() @@ -32,56 +36,35 @@ class YtDlpInfo: """Возвращает title видео""" return self.info.get("title", "unknown") - def get_video_url(self, resolution: Optional[str] = None) -> Optional[MediaContent]: - """ - Возвращает ссылку на видеопоток с указанным разрешением - resolution: например '1080p', '720p', '480p' - """ - formats = self.info.get("formats", []) - video_formats = [ - f - for f in formats - if f.get("vcodec") != "none" # есть видео - and f.get("acodec") == "none" # без аудио + def download_video(self, size: Literal[240, 360, 480, 720, 1080, 1440, 2160]): + command = [ + "yt-dlp", + "-f", + f"bestvideo[height<={size}]", + "--no-part", + "--quiet", + "--no-warnings", + "-o", + "-", + self.url ] - if resolution: - # ищем точное соответствие разрешению - for f in video_formats: - if f.get("format_note") == resolution: - return MediaContent( - f.get("url", ""), - f.get("http_headers", {}), - f.get("downloader_options", {}).get("http_chunk_size", 1024**2), - ) - # если разрешение не указано или не найдено — берем лучший - if video_formats: - # сортируем по height - video_formats.sort(key=lambda x: x.get("height") or 0, reverse=True) - return MediaContent( - video_formats[0].get("url", ""), - video_formats[0].get("http_headers", {}), - video_formats[0] - .get("downloader_options", {}) - .get("http_chunk_size", 1024**2), - ) - return None - - def get_audio_url(self) -> Optional[MediaContent]: - """Возвращает ссылку на аудиопоток""" - formats = self.info.get("formats", []) - audio_formats = [ - f - for f in formats - if f.get("vcodec") == "none" and f.get("acodec") != "none" - ] - if audio_formats: - # берем наилучшее качество - audio_formats.sort(key=lambda x: x.get("abr") or 0, reverse=True) - return MediaContent( - audio_formats[0].get("url", ""), - audio_formats[0].get("http_headers", {}), - audio_formats[0] - .get("downloader_options", {}) - .get("http_chunk_size", 1024**2), - ) - return None + print("Start processing") + process = subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, bufsize=0) + print("Write filename to upload backend") + self.backend.start(self.title + ".mp4") + print("Start write chunk to upload backend") + chunk_size = 1024 ** 2 + length = 0 + while True: + chunk = process.stdout.read(chunk_size) + if not chunk: + break + length += chunk_size + print("Write chunk to backend", length) + self.backend.upload_chunk(chunk) + print("End writing to backend") + ret = process.wait() + print("Check ret status") + if ret != 0: + self.backend.abort() + raise RuntimeError(f"yt-dlp failed, status code: {ret}") diff --git a/test.py b/test.py index 8ffca01..2a974fc 100644 --- a/test.py +++ b/test.py @@ -1,19 +1,21 @@ from app.utils.downloader import HttpStreamingDownloader from app.utils.uploader import DiskChunkUploadBackend -from app.utils.youtube import YtDlpInfo +from app.utils.youtube import YtDlpManager def download(url: str): upload_backend = DiskChunkUploadBackend("trash_holder") - downloader = HttpStreamingDownloader(upload_backend) - youtube = YtDlpInfo(url) - video = youtube.get_video_url("480p") - video_name = youtube.title + ".mp4" - audio = youtube.get_audio_url() - audio_name = youtube.title + ".m4a" + youtube = YtDlpManager(url, upload_backend) + youtube.download_video(360) + # downloader = HttpStreamingDownloader(upload_backend) + # youtube = YtDlpInfo(url) + # video = youtube.get_video_url("480p") + # video_name = youtube.title + ".mp4" + # audio = youtube.get_audio_url() + # audio_name = youtube.title + ".m4a" - downloader.download(video.url, video_name, video.headers, video.chunk_size) - downloader.download(audio.url, audio_name, audio.headers, audio.chunk_size) + # downloader.download(video.url, video_name, video.headers, video.chunk_size) + # downloader.download(audio.url, audio_name, audio.headers, audio.chunk_size) def main():