Add more options to args

1 files changed, 139 insertions, 195 deletions

diff --git a/jamos b/jamos
index 78c377d..8442eae 100755
--- a/jamos
+++ b/jamos
@@ -1,5 +1,6 @@
 #!/usr/bin/env python3
 import argparse
+import collections
 import datetime
 import glob
 import json
@@ -13,6 +14,8 @@ import youtube_dl
 from ytmusicapi import YTMusic
 # import musicpd
 
+# TODO: set this to false to begin with. for now it's always gonna print
+# debug information
 DEBUG = True
 
 # Includes the format characters
@@ -32,129 +35,106 @@ How the albums will be formatted:
     }
 }
 """
-ALBUMS = {}
-
-
-def cleanup_metadata_files(music_directory):
-    files = glob.glob(os.path.join(music_directory, '*.json'))
-
-    for file in files:
-        os.remove(file)
-
-
-def format_youtube_date(date):
-    default = "Unknown Year"
-    try:
-        fmt = "%Y%m%d"
-        d = datetime.datetime.strptime(date, fmt)
-        return d.year
-    except Exception as ex:
-        print(ex)
-
-    return default
-
-
-def get_all_files(directory):
-    things = glob.glob(os.path.join(directory, '*.mp3'))
-
-    files = []
-
-    for thing in things:
-        if os.path.isfile(thing):
-            files.append(thing)
-
-    return files
+ALBUMS = collections.defaultdict(lambda: None)
 
 
 def get_command_line_options():
     parser = argparse.ArgumentParser(
         description="Download songs from YouTube Music")
 
-    parser.add_argument("url",
-                        metavar="https://music.youtube.com/playlist?list=1234",
-                        nargs='?',
-                        default=None,
+    # NOTE: I'm not sure if this param is necessary anymore. If playlist
+    # is public/unlisted, then it shouldn't need the cookies.
+    parser.add_argument("--cookies",
+                        metavar="cookiefile.txt",
+                        # TODO: store this in a config file
+                        default=os.path.join(os.path.expanduser("~"),
+                                             "cookies.txt"),
                         type=str,
-                        help="Playlist or Song URL to download.")
+                        help="Cookie file to use.")
 
-    parser.add_argument("-c",
-                        "--cookies",
-                        metavar="cookiefile.txt",
+    parser.add_argument("--debug_location",
+                        metavar="~/Music/jamos_debug",
+                        default=None,
                         type=str,
                         help="Cookie file to use.")
 
+    # NOTE: unlike cookies, this is required to get the library
     parser.add_argument("--headers",
                         metavar="header_path",
+                        # TODO: store this in a config file
+                        default='~/.config/jamos/headers.json',
                         type=str,
                         help="Header file to use.")
 
     parser.add_argument("-o",
                         "--output",
                         metavar="output_directory",
+                        # TODO: store it in a config file
+                        default=os.path.join(os.path.expanduser("~"), "Music"),
                         type=str,
                         help="Output directory to use")
 
+    parser.add_argument("--playlist_id",
+                        metavar='PLFsQleAWXsj_4yDeebiIADdH5FMayBiJo',
+                        nargs='?',
+                        type=str,
+                        help="""Playlist id to download. Playlist MUST
+                        be public or unlisted, but NOT private""")
+
     parser.add_argument("-r",
-                        "--retryFile",
+                        "--retry_file",
                         metavar="jamos_failed_urls.txt",
                         default=None,
                         type=str,
-                        help="Output directory to use")
-
-    args = parser.parse_args()
-
-    if args.retryFile and args.url:
-        print("Cannot have url and retry flag!\n")
-        parser.print_help()
-        raise Exception()
+                        help="File to store songs that were not completed")
 
-    if not args.retryFile and not args.url:
-        print("Must pass either a url or a retry file!\n")
-        parser.print_help()
-        raise Exception()
-
-    return args
-
-
-def get_song_metadata_from_json(json, counter):
-    metadata = {
-        'title': 'unknownsong',
-        'artist': 'unknownartist',
-        'album': 'unknownalbum',
-        'year': 1999,
-        'jamos_filename': 'jamos_unknwon_file_number_{}.mp3'.format(counter),
-    }
+    parser.add_argument("--song_id",
+                        metavar="lYBUbBu4W08",
+                        nargs='?',
+                        type=str,
+                        help="Song id to download.")
+
+    parser.add_argument("--song_limit",
+                        metavar="5000",
+                        default=5000,
+                        type=int,
+                        help="Number of songs to download")
+
+    parser.add_argument("-t",
+                        "--thumbnail_folder",
+                        metavar="~/Music/jamos_thumbnails",
+                        # TODO: store it in a config file
+                        default=os.path.join(os.path.expanduser("~"), "Music",
+                                             "jamos_thumbnails"),
+                        type=str,
+                        help="Directory to save thumbnails for albums")
 
-    try:
-        if ('artist' in json.keys()) and (json['artist'] is not None):
-            metadata['artist'] = json['artist']
-            if len(metadata['artist'].split(',')) > 1:
-                # If there are multiple artists, pick the first one
-                # NOTE: This will break if the artist has a comma in their name
-                metadata['artist'] = metadata['artist'].split(',')[0]
+    parser.add_argument("-v",
+                        "--verbose",
+                        action='store_true',
+                        help="Directory to save thumbnails for albums")
 
-        if ('album' in json.keys()) and (json['album'] is not None):
-            metadata['album'] = json['album']
+    args = parser.parse_args()
 
-        if ('title' in json.keys()) and (json['title'] is not None):
-            metadata['title'] = json['title']
+    # Validate args
+    if args.verbose:
+        global DEBUG
+        DEBUG = True
 
-        if ('release_date' in json.keys() and
-                json['release_date'] is not None):
-            metadata['year'] = format_youtube_date(json['release_date'])
+    if args.song_id and args.playlist_id:
+        parser.error('--song_id and --playlist_id may not be passed together!')
 
-        artist_for_filename = metadata['artist'].replace(' ', '_').lower()
-        title_for_filename = metadata['title'].replace(' ', '_').lower()
-        metadata['jamos_filename'] = '{}_{}.mp3'.format(artist_for_filename,
-                                                        title_for_filename)
+    if args.song_id and args.retry_file:
+        parser.error('--song_id and --retry_file may not be passed together!')
 
-    except Exception as ex:
-        print(ex)
+    if args.playlist_id and args.retry_file:
+        parser.error(
+            "--playlist_id and --retry_file may not be passed together!")
 
-    return metadata
+    return args
 
 
-def create_downloader(music_directory, cookies):
+def create_downloader(music_directory, cookies=None):
     audio_options = {
         'format': 'mp3/bestaudio/best',
         'cookiefile': cookies,
@@ -174,49 +154,6 @@ def create_downloader(music_directory, cookies):
     return youtube_dl.YoutubeDL(audio_options)
 
 
-def get_video_urls_in_playlist(playlist_url, ytdl):
-    videos = ytdl.extract_info(playlist_url, download=False)
-
-    urls = []
-    for vid in videos['entries']:
-        if 'webpage_url' in vid.keys() and vid['webpage_url'] is not None:
-            urls.append(vid['webpage_url'])
-
-    return urls
-
-
-def move_file(file, metadata, output_directory):
-    artist = remove_special_characters_for_filename(metadata['artist'])
-    album = remove_special_characters_for_filename(metadata['album'])
-    title = remove_special_characters_for_filename(metadata['title'])
-
-    final_directory = os.path.join(
-        output_directory,
-        artist,
-        album)
-
-    Path(final_directory).mkdir(parents=True, exist_ok=True)
-
-    # TODO: Research converting to mp3 instead of just naming it such.
-    # TODO: Research better file formats over mp3?
-    os.rename(
-        file,
-        os.path.join(final_directory, '{}_{}_{}.mp3'.format(artist,
-                                                            album,
-                                                            title)))
-
-
-def save_urls_from_playlist_to_file(filename, urls):
-    try:
-        f = open(filename, "a")
-        for url in urls:
-            f.writelines(url + '\n')
-        f.close()
-    except Exception as e:
-        print(e)
-        raise e
-
-
 def remove_special_characters_for_filename(filename):
     special_chars = [
         ['-', ' '],
@@ -239,53 +176,52 @@ def remove_special_characters_for_filename(filename):
     return new_name.lower()
 
 
-def write_metadata_to_song_file(filename, metadata):
-    file = music_tag.load_file(filename)
-
-    file['name'] = metadata['title']
-    file['artist'] = metadata['artist']
-    file['album'] = metadata['album']
-    file['year'] = metadata['year']
-
-    file.save()
-
-
 def create_youtube_music_api_object(header_path):
     return YTMusic(header_path)
 
 
-def save_album(app, album_id):
+def save_album(app, args, album_id):
     raw_album = app.get_album(album_id)
-    # TODO: Add thumbnail_url
-    ALBUMS[album_id] = {
+
+    # TODO: Add get_thumbnail function
+    thumbnail_url = raw_album['thumbnails'][-1]['url']
+
+    parsed_album = {
         'name': raw_album['title'],
         'tracks': [{'title': song['title'], 'id': song['videoId'],
                     'track_num': index + 1}
                    for index, song in enumerate(raw_album['tracks'])],
         'year': raw_album['year'],
+        'thumbnail_url': thumbnail_url,
+        'thumbnail_filepath': None
     }
 
-    # TODO: Add get_thumbnail
-    thumbnail_url = raw_album['thumbnails'][-1]['url']
-    r = requests.get(thumbnail_url, stream=True)
-    filename = './thumbnails/' + album_id + '.jpg'
+    try:
+        r = requests.get(thumbnail_url, stream=True)
+        filename = os.path.join(args.thumbnail_folder, album_id + '.jpg')
 
-    if not os.path.isfile(filename):
-        # Check if the image was retrieved successfully
-        if r.status_code == 200:
-            # Set decode_content value to True,
-            # otherwise the downloaded image file's size will be zero.
-            r.raw.decode_content = True
+        if not os.path.isfile(filename):
+            # Check if the image was retrieved successfully
+            if r.status_code == 200:
+                # Set decode_content value to True,
+                # otherwise the downloaded image file's size will be zero.
+                r.raw.decode_content = True
 
-            # Open a local file with wb ( write binary ) permission.
-            with open(filename, 'wb') as f:
-                shutil.copyfileobj(r.raw, f)
+                # Open a local file with wb ( write binary ) permission.
+                with open(filename, 'wb') as f:
+                    shutil.copyfileobj(r.raw, f)
 
-            print('Image sucessfully Downloaded: ', filename)
-        else:
-            print('Image Couldn\'t be retreived')
+                print('Image sucessfully Downloaded: ', filename)
+            else:
+                print('Image Couldn\'t be retreived')
 
-    ALBUMS[album_id]['thumbnail_filepath'] = filename
+        parsed_album['thumbnail_filepath'] = filename
+    except Exception as ex:
+        # NOTE: This is gonna be a pain to fix if a thumbnail doesn't download
+        print("Could not download thumbnail")
+        print(ex)
+
+    ALBUMS[album_id] = parsed_album
 
 
 def save_debug_data(filename, data, debug=True):
@@ -294,13 +230,15 @@ def save_debug_data(filename, data, debug=True):
             f.write(data)
 
 
-def parse_songs(app, all_raw_song_data):
+# TODO: Good place for some fancy yield work?
+def parse_songs(app, args, all_raw_song_data):
     all_parsed_songs = []
     for raw_song in all_raw_song_data:
         song_id = raw_song['videoId']
+
         album_id = raw_song['album']['id']
-        if album_id not in ALBUMS.keys():
-            save_album(app, album_id)
+        if not ALBUMS[album_id]:
+            save_album(app, args, album_id)
 
         album = ALBUMS[album_id]
 
@@ -321,21 +259,25 @@ def parse_songs(app, all_raw_song_data):
         }
         all_parsed_songs.append(parsed_song)
 
-    save_debug_data('parsed_album_data.json', json.dumps(ALBUMS),
-                    debug=DEBUG)
-    save_debug_data('parsed_song_data.json', json.dumps(all_parsed_songs),
-                    debug=DEBUG)
+    if args.debug_location:
+        save_debug_data('parsed_album_data.json', json.dumps(ALBUMS),
+                        debug=DEBUG)
+        save_debug_data('parsed_song_data.json', json.dumps(all_parsed_songs),
+                        debug=DEBUG)
 
     return all_parsed_songs
 
 
-def get_library_songs(app, song_limit, order='a_to_z'):
-    all_raw_song_data = app.get_library_songs(limit=song_limit, order=order)
+def get_library_songs(app, args, order='a_to_z'):
+    all_raw_song_data = app.get_library_songs(limit=args.song_limit,
+                                              order=order)
 
-    save_debug_data('raw_song_data.json', json.dumps(all_raw_song_data),
-                    debug=DEBUG)
+    if args.debug_location:
+        path = os.path.join(args.debug_location, 'raw_song_data.json')
+        save_debug_data(path, json.dumps(all_raw_song_data),
+                        debug=DEBUG)
 
-    all_parsed_songs = parse_songs(app, all_raw_song_data)
+    all_parsed_songs = parse_songs(app, args, all_raw_song_data)
 
     return all_parsed_songs
 
@@ -402,14 +344,14 @@ def retry_downloading_songs():
     #         save_urls_from_playlist_to_file(
     #             os.path.join(music_directory, "jamos_failed_urls.txt"),
     #             failed_urls)
-    #     elif args.retryFile:
+    #     elif args.retry_file:
     #        # Just because we don't have any failed urls in this run, doesn't
     #        # mean that we can get rid of the retry file. We'll only remove it
     #         # if it's been explicitly tried and we have no failed urls.
 
     #        # We've successfully downloaded all of the previously failed urls.
     #         # Delete the file
-    #         os.remove(args.retryFile)
+    #         os.remove(args.retry_file)
     # except Exception as ex:
     #     print(ex)
     #     print("Saving failed urls to file failed! Printing failed urls:")
@@ -426,11 +368,11 @@ def handle_downloading_playlist():
     pass
 
 
-def sanitize_for_filename(filename):
-    new_filename = filename.replace(' ', '_')
-    new_filename = (''.join(
-            [s for s in new_filename if s.isalnum() or s == '_'])).lower()
-    return new_filename
+# def sanitize_for_filename(filename):
+#     new_filename = filename.replace(' ', '_')
+#     new_filename = (''.join(
+#             [s for s in new_filename if s.isalnum() or s == '_'])).lower()
+#     return new_filename
 
 
 if __name__ == "__main__":
@@ -440,18 +382,13 @@ if __name__ == "__main__":
         print(ex)
         sys.exit()
 
-    music_directory = args.output or os.path.join(os.path.expanduser("~"),
-                                                  "Music")
-    cookies = args.cookies or os.path.join(os.path.expanduser("~"),
-                                           "cookies.txt")
+    music_directory = args.output
 
     app = create_youtube_music_api_object(args.headers)
 
-    all_songs = get_library_songs(app, 5000)
+    all_songs = get_library_songs(app, args)
 
-    # From some testing, if your playlist is public, you don't have to use a
-    # cookie file. Youtube-dl doesn't break or throw if the file doesn't exist.
-    ytdl = create_downloader(music_directory, cookies)
+    ytdl = create_downloader(music_directory)
 
     for song in all_songs:
         # download song
@@ -472,8 +409,10 @@ if __name__ == "__main__":
 
         # include album cover
         try:
-            with open(song['thumbnail_filepath'], 'rb') as img_in:
-                file['artwork'] = img_in.read()
+            # NOTE: this will be at least None
+            if song['thumbnail_filepath']:
+                with open(song['thumbnail_filepath'], 'rb') as img_in:
+                    file['artwork'] = img_in.read()
         except Exception as ex:
             print(ex)
 
@@ -483,16 +422,21 @@ if __name__ == "__main__":
         # move song to music_directory/artist/album/artist_album_title.ext
         # excluding non filesafe characters
 
-        artist_for_filename = sanitize_for_filename(
-            song['artists'][0].replace('$', 's'))  # exception for $uicideboy$
-        album_for_filename = sanitize_for_filename(song['album'])
-        title_for_filename = sanitize_for_filename(song['title'])
+        artist_for_filename = remove_special_characters_for_filename(
+            song['artists'][0])
+        album_for_filename = remove_special_characters_for_filename(
+            song['album'])
+        title_for_filename = remove_special_characters_for_filename(
+            song['title'])
         song_output_dir = os.path.join(music_directory,
                                        artist_for_filename,
                                        album_for_filename)
+
         Path(song_output_dir).mkdir(parents=True, exist_ok=True)
+
         new_filename = '{}_{}_{}.mp3'.format(artist_for_filename,
                                              album_for_filename,
                                              title_for_filename)
-        shutil.move(filename, os.path.join(song_output_dir, new_filename))
 
+        shutil.move(filename, os.path.join(song_output_dir, new_filename))
+    print("done")