flac support and an argument for choosing when to embed cover art

bcao.py

@@ -6,19 +6,20 @@
 # output: it organises it, adds cover art, puts it in the right place...
 
 import argparse
-import base64
 import os
 import re
 import sys
 import tempfile
 import shutil
 from os import path
+from base64 import b64encode
 from zipfile import ZipFile
 from pathlib import Path
 from typing import Optional, List, Dict
 
 import mutagen
-from mutagen.flac import Picture
+from mutagen.flac import Picture, FLAC
+from mutagen.oggvorbis import OggVorbis
 from mutagen import id3
 
 from PIL import Image
@@ -37,24 +38,39 @@ def get_tag(mut_song: mutagen.FileType, tag: str) -> str:
 	tag = tag.replace("_", "")
 	return sanitise(mut_song[tag][0] if type(mut_song[tag]) is list else mut_song[tag])
 
+def has_cover(mut_song: mutagen.FileType):
+	if type(mut_song) is OggVorbis:
+		return "metadata_block_picture" in mut_song and len(mut_song["metadata_block_picture"]) != 0
+	if type(mut_song) is FLAC:
+		return len(mut_song.pictures) != 0
+	raise Exception(f"what is {type(mut_song)}")
+
 def sanitise(in_str: str) -> str:
 	if args.sanitise:
 		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')
+def parse_add_cover_images(value: str) -> str:
+	if value not in ["n", "a", "w"]:
+		raise argparse.ArgumentTypeError("Must be one of n, a, w")
+	return value
+
+parser = argparse.ArgumentParser(usage='%(prog)s zip [options]',
+	description="Extracts the given zip file downloaded from Bandcamp and organises it.")
+parser.add_argument('zip', help='The zip file to use.')
+parser.add_argument('-c', '--add-cover-images', dest='process_cover', default='w', type=parse_add_cover_images,
+	help="When to embed cover art into songs. Options: [n]ever, [a]lways, [w]hen necessary. Default: %(default)s")
 parser.add_argument('-d', '--destination', dest='destination', default='/home/lynne/Music/Music/',
-	help="The directory to organise the music into. Default: /home/lynne/Music/Music/")
+	help="The directory to organise the music into. Default: %(default)s")
 parser.add_argument('-q', '--quiet', dest='quiet', action='store_true',
 	help='Disable non-error output and assume default artist name.')
 parser.add_argument('-u', '--unsanitised', dest='sanitise', action='store_false',
 	help="Don't replace NTFS-unsafe characters with underscores. Not recommended.")
 parser.add_argument('-t', '--threshold', dest='threshold', nargs=1, default=300,
-	help="Maximum acceptable cover art file size in kilobytes. Default: 300")
+	help="Maximum acceptable file size for cover art, in kilobytes. Default: %(default)s")
 
 args = parser.parse_args()
-# convert args.treshold to bytes
+# convert args.threshold to bytes
 args.threshold *= 1024
 
 if not path.exists(args.zip):
@@ -68,7 +84,7 @@ song_names: List[str] = []
 
 with ZipFile(args.zip, 'r') as zip_file:
 	for file in zip_file.namelist():
-		if re.match(r"^(.+ - ){2}\d{2,} .+\.(ogg|flac|alac|aiff|wav|mp3|opus|m4a|aac|oga)$", file):
+		if re.match(r"^(.+ - ){2}\d{2,} .+\.(ogg|flac|alac|aiff|wav|mp3|opus|m4a|aac)$", file):
 			# bandcamp zips contains songs with names formatted like "Album - Artist - 01 Song.mp3"
 			# for example, "King Crimson - In the Wake of Poseidon - 02 Pictures of a City.ogg"
 			# this regex should match only on those, and cut out (hopefully) all of the bonus material stuff, which shouldn't
@@ -84,8 +100,9 @@ with ZipFile(args.zip, 'r') as zip_file:
 # we'll need this to know what metadata format we should write
 song_format: str = path.splitext(song_names[0])[1][1:]
 
-log("Resizing album art to embed in songs...")
-with Image.open(str(Path(tmp, cover))) as image:
+if args.process_cover != 'n':
+	log("Resizing album art to embed in songs...")
+	with Image.open(str(Path(tmp, cover))) as image:
 		temp_cover: Path = Path(tmp, "cover-lq.jpg")
 
 		if image.mode in ["RGBA", "P"]:
@@ -109,13 +126,12 @@ with Image.open(str(Path(tmp, cover))) as image:
 				# something very bad has happened here
 				die("Failed to resize image")
 
-
-# read the image file to get the file's raw data
-with open(temp_cover, 'r+b') as cover_file:
+	# read the image file to get its raw data
+	with open(temp_cover, 'r+b') as cover_file:
 		data = cover_file.read()
 
-with Image.open(temp_cover) as image:
-	if song_format == "ogg":
+	with Image.open(temp_cover) as image:
+		if song_format == "ogg" or song_format == "flac":
 			# i hate this
 			embed_cover = Picture()
 			embed_cover.data = data
@@ -124,7 +140,6 @@ with Image.open(temp_cover) as image:
 			embed_cover.width = image.size[0]
 			embed_cover.height = image.size[1]
 			embed_cover.depth = image.bits
-		embed_cover = base64.b64encode(embed_cover.write()).decode("ascii")
 		else:
 			log(f"Format {song_format} is not fully supported - cover images will not be modified", 1)
 
@@ -146,9 +161,15 @@ for song in song_names:
 	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}"
 
+	if args.process_cover == 'a' or (args.process_cover == 'w' and has_cover(m) is False):
+		log("Embedding cover art...")
 		# embed cover art
 		if song_format == "ogg":
-		m["metadata_block_picture"] = [embed_cover]
+			m["metadata_block_picture"] = [b64encode(embed_cover.write()).decode("ascii")]
+			m.save()
+		elif song_format == "flac":
+			m.clear_pictures()
+			m.add_picture(embed_cover)
 			m.save()
 
 # remove duplicate artists