♻️ create handleItemDownload

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