From 801b7369c487be2f96f0a0489dcfb66fc1486047 Mon Sep 17 00:00:00 2001 From: Lynnesbian Date: Sat, 17 Oct 2020 16:23:55 +1000 Subject: [PATCH] cleaner code, type annotations, and it even runs (slightly) faster! =u= --- bcao.py | 116 ++++++++++++++++++++++++++++++++------------------------ 1 file changed, 67 insertions(+), 49 deletions(-) diff --git a/bcao.py b/bcao.py index 20f2337..68c7c82 100755 --- a/bcao.py +++ b/bcao.py @@ -9,10 +9,13 @@ import argparse import base64 import os import re -import subprocess +import sys import tempfile import shutil -from zipfile import ZipFile, BadZipFile +from os import path +from zipfile import ZipFile +from pathlib import Path +from typing import Optional, List, Dict import mutagen from mutagen.flac import Picture @@ -26,24 +29,18 @@ def log(message: str, importance: int = 0): def die(message: str, code: int = 1): print(message) - exit(code) + sys.exit(code) -def get_tag(m: mutagen.FileType, tag: str): - if tag == "title": - return sanitise(m['title'][0]) - elif tag == "track": - return int(m['tracknumber'][0]) - elif tag == "album": - return sanitise(m['album'][0]) - else: - # may as well try - return sanitise(m[tag]) +def get_tag(mut_song: mutagen.FileType, tag: str) -> str: + if tag == "track": + tag = "tracknumber" + tag = tag.replace("_", "") + return sanitise(mut_song[tag][0] if type(mut_song[tag]) is list else mut_song[tag]) -def sanitise(input: str): +def sanitise(in_str: str) -> str: if args.sanitise: - return re.sub(r"[?\\/:|*\"<>]", "_", input) - else: - return input + return re.sub(r"[?\\/:|*\"<>]", "_", in_str) + return in_str parser = argparse.ArgumentParser(description="Extracts the given zip file downloaded from Bandcamp and organises it.") parser.add_argument('zip', help='The zip file to use') @@ -57,15 +54,17 @@ parser.add_argument('-t', '--threshold', dest='threshold', nargs=1, default=300, help="Maximum acceptable cover art file size in kilobytes. Default: 300") args = parser.parse_args() +# convert args.treshold to bytes +args.threshold *= 1024 -if not os.path.exists(args.zip): +if not path.exists(args.zip): die(f"Couldn't find {args.zip}.", 2) log("Extracting...") -tmpd = tempfile.TemporaryDirectory() -tmp = tmpd.name -cover = None -song_names = [] +tmp_dir: tempfile.TemporaryDirectory = tempfile.TemporaryDirectory() +tmp: str = tmp_dir.name +cover: Optional[str] = None +song_names: List[str] = [] with ZipFile(args.zip, 'r') as zip_file: for file in zip_file.namelist(): @@ -83,12 +82,11 @@ with ZipFile(args.zip, 'r') as zip_file: # save the format of the songs (ogg, mp3, etc) # we'll need this to know what metadata format we should write -song_format = os.path.splitext(song_names[0])[1][1:] +song_format: str = path.splitext(song_names[0])[1][1:] log("Resizing album art to embed in songs...") - -with Image.open(os.path.join(tmp, cover)) as image: - temp_cover = os.path.join(tmp, "cover-lq.jpg") +with Image.open(str(Path(tmp, cover))) as image: + temp_cover: Path = Path(tmp, "cover-lq.jpg") if image.mode in ["RGBA", "P"]: # remove alpha channel @@ -97,11 +95,18 @@ with Image.open(os.path.join(tmp, cover)) as image: image.save(temp_cover, quality=85, optimize=True) image_smol = image - # keep shrinking the image by 90% until it's less than {args.threshold} kilobytes - while os.path.getsize(temp_cover) / 1024 > args.threshold: - image_smol = image_smol.resize((round(image_smol.size[0] * 0.9), round(image_smol.size[1] * 0.9))) + while path.getsize(temp_cover) > args.threshold: + # keep shrinking the image by 90% until it's less than {args.threshold} kilobytes + ratio = 0.9 + + if path.getsize(temp_cover) > args.threshold * 2: + # if the file size of the cover is more than double the threshold, resize the cover image size by 80% instead + ratio = 0.8 + + image_smol = image_smol.resize([round(n * ratio) for n in image_smol.size]) image_smol.save(temp_cover, quality=85, optimize=True) if image_smol.size[0] == 10: + # something very bad has happened here die("Failed to resize image") @@ -123,50 +128,63 @@ with Image.open(temp_cover) as image: else: log(f"Format {song_format} is not fully supported - cover images will not be modified", 1) -artists = [] -album = None -songs = {} +artists: List[str] = [] +album: Optional[str] = None +songs: Dict[str, str] = {} zeroes = min(len(song_names), 2) +first_loop: bool = True + for song in song_names: - ext = os.path.splitext(song)[1:] - m = mutagen.File(os.path.join(tmp, song)) - # add the song's artist to the list, if it hasn't been seen yet - [artists.append(sanitise(artist)) for artist in m['artist'] if artist not in artists] + m = mutagen.File(Path(tmp, song)) + if first_loop: + # the first item in the artists list should be the album artist + artists.append(get_tag(m, "album_artist")) + album = get_tag(m, "album") + first_loop = False + + # add the song's artist(s) to the list + map(lambda x: artists.append(sanitise(x)), m['artist']) songs[song] = f"{str(get_tag(m, 'track')).zfill(zeroes)} {get_tag(m, 'title')}.{song_format}" - album = get_tag(m, "album") + # embed cover art if song_format == "ogg": m["metadata_block_picture"] = [embed_cover] m.save() +# remove duplicate artists +artists = list(dict.fromkeys(artists)) + if len(artists) > 1 and "Various Artists" not in artists: artists.append("Various Artists") -artist = None +artist: Optional[str] = None while artist is None: log("Artist directory:") - for i in range(len(artists)): - log(f"{i+1}) {artists[i]}") + for i, artist_name in enumerate(artists): + log(f"{i+1}) {artist_name}") log(f"{len(artists) + 1}) Custom...") + choice = "1" if args.quiet else input("> ") if choice.isdecimal(): - if int(choice) == len(artists) + 1: - log("Enter the name to use.") + choice = int(choice) + if choice == len(artists) + 1: + log("Enter the name to use:") + artist = input("> ") else: try: - artist = artists[int(choice) - 1] + artist = artists[choice - 1] except KeyError: log(f"Please choose a number between 1 and {len(artists) + 1}.") else: log(f"Please choose a number between 1 and {len(artists) + 1}") -destination = os.path.join(args.destination, artist, album) -log(f"Moving files to {destination}...") +destination: Path = Path(args.destination, artist, album) +log(f"Moving files to \"{destination}\"...") os.makedirs(destination, exist_ok=True) + for source_name, dest_name in songs.items(): - shutil.move(os.path.join(tmp, source_name), os.path.join(destination, dest_name)) -shutil.move(os.path.join(tmp, cover), os.path.join(destination, cover)) + shutil.move(Path(tmp, source_name), Path(destination, dest_name)) +shutil.move(Path(tmp, cover), Path(destination, cover)) -tmpd.cleanup() +tmp_dir.cleanup() log("Done!") -