cleaner code, type annotations, and it even runs (slightly) faster! =u=

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!")
-