Use ytmusic API to query album ID

1 files changed, 26 insertions, 4 deletions

diff --git a/yt_music_scraper/main.py b/yt_music_scraper/main.py
index 52e344f..3e8a917 100644
--- a/yt_music_scraper/main.py
+++ b/yt_music_scraper/main.py
@@ -1,4 +1,7 @@
 from yt_dlp import YoutubeDL
+import urllib.parse
+from ytmusicapi import YTMusic
+import re
 import json
 import argparse
 import sys
@@ -10,12 +13,18 @@ LOGLEVEL = os.environ.get('LOG_LEVEL', 'INFO').upper()
 logger.setLevel(LOGLEVEL)
 logger.debug("Starting...")
 
+def album_info(data):
+
+    logger.info(json.dumps(data['playlist'] if 'playlist' in data else "N/A", indent=4))
+    logger.info(json.dumps(data['playlist_id'] if 'playlist_id' in data else "N/A", indent=4))
+    return data['playlist_id'] if 'playlist_id' in data else "N/A"
+
 def vid_info(data):
 
     logger.info(json.dumps(data['title'] if 'title' in data else "N/A", indent=4))
     logger.info(json.dumps(data['album'] if 'album' in data else "N/A", indent=4))
     logger.info(json.dumps(data['artists'] if 'artists' in data else "N/A", indent=4))
-    return (data['album'] if 'album' in data else "N/A", data['artists'] if 'artists' in data else "N/A")
+    return data['album'] if 'album' in data else "N/A"
 
 ydl_opts = {
     'format': 'bestaudio/best',
@@ -25,6 +34,7 @@ ydl_opts = {
         'preferredquality': 'best',
     }],
     'logger': logger,
+    'output': '%(playlist)s/%(playlist_index)s - %(title)s.%(ext)s',
 }
 
 
@@ -36,14 +46,26 @@ if __name__ == "__main__":
     args = parser.parse_args()
 
     logger.debug("Create YT object...")
+    albums = []
     yt = YoutubeDL(ydl_opts)
     data = yt.extract_info(args.playlist, download=False)
-    with open("data.json", "w") as f:
-        f.write(json.dumps(data, indent=4))
     i = 0
     if data:
-        albums = []
+        albums = set()
         for datum in data["entries"]:
             logger.debug(json.dumps(datum, indent=4))
+            albums.add(vid_info(datum))
             logger.info(f"ALBUM NAME: {vid_info(datum)}")
 
+        logger.info(f"albums: {albums}")
+        re = re.compile("OLAK.*")
+        ytmusic = YTMusic()
+        album_ids = []
+        for album in list(albums):
+            results = ytmusic.search(albums.pop())
+            for item in results:
+                if item['resultType'] == 'album':
+                    album_ids.append(item['playlistId'])
+                    break
+        logger.info(f"Album IDs: {album_ids}")
+