musick/musick.py

import os
from mutagen.easyid3 import EasyID3
from mutagen.id3 import ID3, APIC, TIT2, TPE1, TALB
import getopt
import sys
import shutil

total_proc = 0
total_succ = 0
total_error = 0
total_cover = 0

input_dir = None
output_dir = None
cover_dir = None
cover_arts = []

UNKNOWN_ARTIST_NAME = "Unknown Artist"
UNKNOWN_ALBUM_NAME = "Unknown Album"

WILDCARD_MATCH_SCORE = 1
ARTIST_MATCH_SCORE = 10
ALBUM_MATCH_SCORE = 100
TITLE_MATCH_SCORE = 1000

# to get a > 0 score, each field either must match or the cover tuple is wildcare (None)
def calc_cover_score(artist: str, album: str, title: str, tup : tuple):
    score = 0

    if (tup[0] == artist):
        score += ARTIST_MATCH_SCORE
    elif (tup[0] == ""):
        score += WILDCARD_MATCH_SCORE
    else:
        # mismatched artist, discard
        return 0
    
    if (tup[1] == album):
        score += ALBUM_MATCH_SCORE
    elif (tup[1] == ""):
        score += WILDCARD_MATCH_SCORE
    else:
        # mismatched album, discard
        return 0 

    if (tup[2] == title):
        score += TITLE_MATCH_SCORE
    elif (tup[2] == ""):
        score += WILDCARD_MATCH_SCORE
    else:
        # mismatched title, discard
        return 0

    return score

# returns the idx into cover_arts array.
# < 0 means no match
def find_best_cover(artist: str, album: str, title: str):
    best_idx = -1
    best_score = 0
    for i in range(0, len(cover_arts)):
        score = calc_cover_score(artist, album, title, cover_arts[i])
        if score > best_score:
            best_idx = i
            best_score = score

    return best_idx


def parse_triple(f: str):
    triple = os.path.splitext(os.path.basename(f))[0].split("_")
    album = ""
    title = ""
    artist = triple[0]
    if len(triple) >= 2:
        album = triple[1]
    if len(triple) >= 3:
        title = triple[2]
    return artist, album, title

def load_cover_arts():
    global cover_arts

    if cover_dir == None:
        return

    cdir = cover_dir
    files = os.listdir(cdir)
    for f in files:
        try:
            artist, album, title = parse_triple(f)
            f = os.path.join(cover_dir, f)
            print("Processing cover art - " + f)
            with open(f,'rb') as img:
                cover_arts.append((artist, album, title, img.read(), os.path.basename(f)))
            print("   Loading cover art - " + f + " Artist: " + artist + " Album: " + album + " Title: " + title)
        except Exception as e:
            print("   Skipping due to Exception - " + str(e))
        print("")
        

def usage():
    print("Usage: python musik.py [options]\n\n\
    options:\n\
        -h : show usage.\n\
        -i input : input directory\n\
        -o output : output directory\n\
        -c cover : cover arts page\n")


no_match = []

def process_file(f: str, odir: str):
    global total_cover
    global total_succ
    global total_error
    global total_proc
    try:
        print("Processing file - " + f)
        total_proc += 1

        artist, album, title = parse_triple(f)

        # id3 = ID3(f)

        # album = str(id3.get("TALB") if id3.get("TALB") != None else "")
        if album == "":
            album = UNKNOWN_ALBUM_NAME

        # artist = str(id3.get("TPE1") if id3.get("TPE1") != None else "")
        if artist == "":
            artist = UNKNOWN_ARTIST_NAME
        
        # title = str(id3.get("TIT2") if id3.get("TIT2") != None else "")

        if title == "":
            raise Exception("No title")
        
        print("    Artist: " + artist + " Album: " + album + " Title: " + title)

        target_dir = os.path.join(odir, artist, album)
        os.makedirs(target_dir, exist_ok = True)

        target_f = os.path.join(target_dir, os.path.basename(f))
        try:
            shutil.copyfile(f, target_f)
        except shutil.SameFileError:
            pass
        print("    Copied file to - " + target_f)

        cover_idx = find_best_cover(artist, album, title)
        
        # write id3 tags of the target file
        id3 = ID3(target_f)
        # write title
        id3.delall("TIT2")
        id3.add(TIT2(text = title))
        # write album
        id3.delall("TALB")
        id3.add(TALB(text = album))
        # write artist
        id3.delall("TPE1")
        id3.add(TPE1(text = artist))

        if cover_idx >= 0:
            # install cover
            id3.delall("APIC")
            id3.add(APIC(3, "image/jpeg", 3, "Front cover", cover_arts[cover_idx][3]))
            total_cover += 1
            print("    Written cover art - " + cover_arts[cover_idx][4])
        else:
            no_match.append(target_f)

        id3.save(v2_version=3)

        total_succ += 1

    except Exception as e:
        total_error += 1
        print("    Skipping due to exception - %s" % str(e))
    
    print("")
    
def process_directory(idir: str, odir: str):
    print("Processing directory - " + idir)
    files = os.listdir(idir)
    for f in files:
        f = os.path.join(idir, f)
        if os.path.isdir(f):
            process_directory(f, odir)
        else:
            process_file(f, odir)


def main():
    global input_dir
    global output_dir
    global cover_dir

    try:
        opts , _ = getopt.getopt(sys.argv[1:], "hi:o:c:")
    except getopt.GetoptError as err:
        print(str(err))
        usage()
        sys.exit(1)

    for o, a in opts:
        if o == "-h":
            usage()
            sys.exit(0)
        elif o == "-i":
            input_dir = a
        elif o == "-o":
            output_dir = a
        elif o == "-c":
            cover_dir = a
        else:
            print("Unrecognized option: " + o)
            sys.exit(1)
    
    if input_dir == None or output_dir == None:
        print("Must specify both -i and -o")
        sys.exit(1)

    load_cover_arts()
    process_directory(input_dir, output_dir)

    print("\nSummary - Total Files: " + str(total_proc) + " Total Errors: " + str(total_error) + " Total Covers: " + str(total_cover))
    for each in no_match:
        print("No match: " + each)

main()