From 672f56597eeaa597db0b566ebb3d542e1f0cea7f Mon Sep 17 00:00:00 2001 From: oskvr37 Date: Thu, 6 Feb 2025 23:11:08 +0100 Subject: [PATCH] =?UTF-8?q?=E2=99=BB=EF=B8=8F=20create=20`handleItemDownlo?= =?UTF-8?q?ad`?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit --- examples/concurrent_download_rich.py | 113 ++++++++++++--------------- 1 file changed, 48 insertions(+), 65 deletions(-) diff --git a/examples/concurrent_download_rich.py b/examples/concurrent_download_rich.py index a1bd9c8..a1ec780 100644 --- a/examples/concurrent_download_rich.py +++ b/examples/concurrent_download_rich.py @@ -5,6 +5,8 @@ This will download tracks and videos. import logging +from typing import Union + from pathlib import Path from requests import Session from concurrent.futures import ThreadPoolExecutor @@ -12,7 +14,6 @@ from concurrent.futures import ThreadPoolExecutor from rich.console import Console from rich.logging import RichHandler from rich.progress import ( - TaskID, BarColumn, Progress, TextColumn, @@ -21,7 +22,6 @@ from rich.progress import ( from tiddl.api import TidalApi from tiddl.download import parseTrackStream, parseVideoStream from tiddl.config import Config -from tiddl.models.api import PlaylistItems from tiddl.models.resource import Track, Video from tiddl.utils import convertFileExtension @@ -29,6 +29,7 @@ from tiddl.utils import convertFileExtension WORKERS_COUNT = 4 PLAYLIST_UUID = "84974059-76af-406a-aede-ece2b78fa372" ALBUM_ID = 103805723 +QUALITY = "HI_RES_LOSSLESS" console = Console() logging.basicConfig( @@ -50,12 +51,28 @@ progress = Progress( ) -def handleTrackDownload(task_id: TaskID, track: Track): - track_stream = api.getTrackStream(track.id, "LOW") - urls, extension = parseTrackStream(track_stream) +def handleItemDownload(item: Union[Track, Video]): + # TODO: check if item is already downloaded - progress.update(task_id, total=len(urls), visible=True) - progress.start_task(task_id) + if isinstance(item, Track): + track_stream = api.getTrackStream(item.id, quality=QUALITY) + urls, extension = parseTrackStream(track_stream) + elif isinstance(item, Video): + video_stream = api.getVideoStream(item.id) + urls = parseVideoStream(video_stream) + extension = ".ts" + else: + raise TypeError( + f"Invalid item type: expected an instance of Track or Video, " + f"received an instance of {type(item).__name__}. " + ) + + task_id = progress.add_task( + description=f"{type(item).__name__} {item.title}", + start=True, + visible=True, + total=len(urls), + ) with Session() as s: stream_data = b"" @@ -63,42 +80,34 @@ def handleTrackDownload(task_id: TaskID, track: Track): for url in urls: req = s.get(url) stream_data += req.content - progress.update(task_id, advance=1) + progress.advance(task_id) - path = Path("tracks") / f"{track.title}{extension}" + path = Path("examples") / "downloads" / f"{item.id}{extension}" path.parent.mkdir(parents=True, exist_ok=True) with path.open("wb") as f: f.write(stream_data) - console.log(track.title) - progress.remove_task(task_id) + if isinstance(item, Track): + if item.audioQuality == "HI_RES_LOSSLESS": + convertFileExtension( + source_file=path, + extension=".flac", + remove_source=True, + is_video=False, + copy_audio=True, # extract flac from m4a container + ) + elif isinstance(item, Video): + convertFileExtension( + source_file=path, + extension=".mp4", + remove_source=True, + is_video=True, + copy_audio=True, + ) -def handleVideoDownload(task_id: TaskID, video: Video): - video_stream = api.getVideoStream(video.id) - urls = parseVideoStream(video_stream) - - progress.update(task_id, total=len(urls), visible=True) - progress.start_task(task_id) - - with Session() as s: - video_data = b"" - - for url in urls: - req = s.get(url) - video_data += req.content - progress.update(task_id, advance=1) - - path = Path("videos") / f"{video.id}.ts" - path.parent.mkdir(parents=True, exist_ok=True) - - with path.open("wb") as f: - f.write(video_data) - - convertFileExtension(path, ".mp4", remove_source=True, is_video=True) - - console.log(video.title) + console.log(item.title) progress.remove_task(task_id) @@ -107,24 +116,8 @@ progress.start() pool = ThreadPoolExecutor(max_workers=WORKERS_COUNT) -def submitTrack(track: Track): - task_id = progress.add_task( - description=track.title, - start=False, - visible=False, - ) - - pool.submit(handleTrackDownload, task_id=task_id, track=track) - - -def submitVideo(video: Video): - task_id = progress.add_task( - description=video.title, - start=False, - visible=False, - ) - - pool.submit(handleVideoDownload, task_id=task_id, video=video) +def submitItem(item: Union[Track, Video]): + pool.submit(handleItemDownload, item=item) # NOTE: these api requests will run one by one, @@ -133,22 +126,12 @@ def submitVideo(video: Video): playlist_items = api.getPlaylistItems(playlist_uuid=PLAYLIST_UUID, limit=10) for item in playlist_items.items: - item = item.item - - if isinstance(item, PlaylistItems.PlaylistTrackItem.PlaylistTrack): - submitTrack(item) - elif isinstance(item, PlaylistItems.PlaylistVideoItem.PlaylistVideo): - submitVideo(item) + submitItem(item.item) album_items = api.getAlbumItems(album_id=ALBUM_ID, limit=5) for item in album_items.items: - item = item.item - - if isinstance(item, Track): - submitTrack(item) - elif isinstance(item, Video): - submitVideo(item) + submitItem(item.item) # cleanup