misc / RQ bulk downloads for TRip

utils/sr_wrapper.py

@@ -1,10 +1,16 @@
 from typing import Optional
+from uuid import uuid4
+from urllib.parse import urlparse
+import hashlib
 import logging
 import os
 import asyncio
-from streamrip.client import TidalClient
-from streamrip.config import Config as StreamripConfig
+import aiohttp
+from streamrip.client import TidalClient  # type: ignore
+from streamrip.config import Config as StreamripConfig  # type: ignore
 from dotenv import load_dotenv
+
+
 load_dotenv()
 
 
@@ -180,8 +186,7 @@ class SRUtil:
             Optional[str]: The stream URL or None if not found.
         """
         if quality not in ["LOSSLESS", "HIGH", "LOW"]:
-            logging.error("Invalid quality requested: %s",
-                          quality)
+            logging.error("Invalid quality requested: %s", quality)
         quality_int: int = int(self.streamrip_config.session.tidal.quality)
         match quality:
             case "HIGH":
@@ -196,12 +201,12 @@ class SRUtil:
         try:
             track = await self.streamrip_client.get_downloadable(
                 track_id=track_id_str, quality=quality_int
-                )
+            )
         except AttributeError:
             await self.streamrip_client.login()
             track = await self.streamrip_client.get_downloadable(
                 track_id=track_id_str, quality=quality_int
-                )            
+            )
         if not track:
             logging.warning("No track found for ID: %s", track_id)
             return None
@@ -210,3 +215,54 @@ class SRUtil:
             logging.warning("No stream URL found for track ID: %s", track_id)
             return None
         return stream_url
+
+    async def get_metadata_by_track_id(self, track_id: int) -> Optional[dict]:
+        if not self.streamrip_client.logged_in:
+            await self.streamrip_client.login()
+        try:
+            metadata = await self.streamrip_client.get_metadata(str(track_id), "track")
+            return {
+                "artist": metadata.get("artist", {}).get("name", "Unknown Artist"),
+                "album": metadata.get("album", {}).get("title", "Unknown Album"),
+                "song": metadata.get("title", uuid4()),
+            }
+        except Exception as e:
+            logging.critical(
+                "Get metadata for %s failed, Exception: %s", track_id, str(e)
+            )
+            return None
+
+    async def download(self, track_id: int, quality: str = "LOSSLESS") -> bool | str:
+        """Download track
+        Args:
+            track_id (int)
+            quality (str): The quality of the stream, default is "LOSSLESS". Other options: HIGH, LOW
+        Returns:
+            bool
+        """
+        if not self.streamrip_client.logged_in:
+            await self.streamrip_client.login()
+        try:
+            track_url = await self.get_stream_url_by_track_id(track_id)
+            if not track_url:
+                return False
+            parsed_url = urlparse(track_url)
+            parsed_url_filename = os.path.basename(parsed_url.path)
+            parsed_url_ext = os.path.splitext(parsed_url_filename)[1]
+            unique = hashlib.sha256(uuid4().bytes).hexdigest()[:16]
+            dl_folder_path = (
+                f"{self.streamrip_config.session.downloads.folder}/{unique}"
+            )
+            dl_path = f"{dl_folder_path}/{track_id}.{parsed_url_ext}"
+            async with aiohttp.ClientSession() as session:
+                async with session.get(
+                    track_url, headers={}, timeout=aiohttp.ClientTimeout(total=60)
+                ) as resp:
+                    resp.raise_for_status()
+                    with open(dl_path, "wb") as f:
+                        async for chunk in resp.content.iter_chunked(1024 * 64):
+                            f.write(chunk)
+                    return dl_path
+        except Exception as e:
+            logging.critical("Error: %s", str(e))
+            return False