bcao/bcao.py

#!/usr/bin/env python3
# BCAO - BandCamp Automatic Organiser
# copyright 2018-2019 @LynnearSoftware@fedi.lynnesbian.space
# Licensed under the GPLv3: https://www.gnu.org/licenses/gpl-3.0.html#content
# input: a .zip from bandcamp
# output: it organises it, adds cover art, puts it in the right place...

import argparse
import base64
import os
import re
import subprocess
import tempfile
import shutil
from zipfile import ZipFile, BadZipFile

import mutagen
from mutagen.flac import Picture
from mutagen import id3

from PIL import Image

def log(message: str, importance: int = 0):
	if not args.quiet or importance > 0:
		print(message)

def die(message: str, code: int = 1):
	print(message)
	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 sanitise(input: str):
	if args.sanitise:
		return re.sub(r"[?\\/:|*\"<>]", "_", input)
	else:
		return input

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')
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/")
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")

args = parser.parse_args()

if not os.path.exists(args.zip):
	die(f"Couldn't find {args.zip}.", 2)

log("Extracting...")
tmpd = tempfile.TemporaryDirectory()
tmp = tmpd.name
cover = None
song_names = []

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):
			# 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
			#  be added to the music folder (since i sync that to my phone, and "making of" videos are cool, but i don't
			#  have the space for it).
			song_names.append(file)
			zip_file.extract(file, tmp)
		elif cover is None and re.match(r"cover\.(jpe?g|png)", file):
			cover = file
			zip_file.extract(file, tmp)

# 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:]

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

	if image.mode in ["RGBA", "P"]:
		# remove alpha channel
		image = image.convert("RGB")

	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)))
		image_smol.save(temp_cover, quality=85, optimize=True)
		if image_smol.size[0] == 10:
			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:
	data = cover_file.read()

with Image.open(temp_cover) as image:
	if song_format == "ogg":
		# i hate this
		embed_cover = Picture()
		embed_cover.data = data
		embed_cover.type = mutagen.id3.PictureType.COVER_FRONT
		embed_cover.mime = "image/jpeg"
		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)

artists = []
album = None
songs = {}
zeroes = min(len(song_names), 2)
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]
	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()

if len(artists) > 1 and "Various Artists" not in artists:
	artists.append("Various Artists")

artist = None
while artist is None:
	log("Artist directory:")
	for i in range(len(artists)):
		log(f"{i+1}) {artists[i]}")
	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.")
		else:
			try:
				artist = artists[int(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}...")
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))

tmpd.cleanup()
log("Done!")