Get/set popularimeter rating ID3 tag for Traktor

traktor_rating.py

# Gets or sets popularimeter ID3 rating tag for Traktor.
# Needs mutagen for ID3 tag manipulation:
#   pip install mutagen
import logging
import os
import sys

import click
import mutagen
import mutagen.id3
from mutagen.id3._frames import POPM
from mutagen.mp4 import MP4FreeForm

logger = logging.getLogger(__name__)

ALLOWED_TYPES = ['mp3', 'flac', 'm4a']


def validate_rating(ctx, param, value):
    if value not in range(1, 6):
        raise click.BadParameter('Rating must be a value between 1 and 5!')
    return value


def set_rating(filename, rating):
    """
    Set Traktor-style popularimeter rating of MP3 file

    :param filename:
    :param rating:
    :return:
    """
    if filename.endswith('.mp3'):
        id3 = mutagen.id3.ID3(filename)
        frame = POPM(email=None, rating=rating * 51, count=0)
        id3.add(frame)
        id3.save()
        print("Written rating {} to file {}".format(rating, filename))

    elif filename.endswith('.flac'):
        audio = mutagen.File(filename)
        audio['rating wmp'] = [str(rating * 51)]
        audio.save()
        print("Written rating {} to file {}".format(rating, filename))

    elif filename.endswith('.m4a'):
        audio = mutagen.File(filename)
        audio.tags['----:com.apple.iTunes:rating wmp'] = (MP4FreeForm(str(rating * 51).encode()))
        audio.save()
        print("Written rating {} to file {}".format(rating, filename))

    elif filename.endswith('.opus'):
        raise Exception("Opus not supported by Traktor - don't know how to set rating.")


def get_rating(filename):
    """
    Get Traktor-style popularimeter rating of MP3 file

    :param filename: file
    :return:
    """
    if filename.endswith('.mp3'):
        audio = mutagen.File(filename)
        if 'POPM:[email protected]' not in audio:
            return None
        rating = mutagen.File(filename)['POPM:[email protected]'].rating
        rating = int(rating / 51)
        return rating

    elif filename.endswith('.flac'):
        audio = mutagen.File(filename)
        if 'rating wmp' not in audio:
            return None
        rating = int(int(audio['rating wmp'][0]) / 51)
        return rating

    elif filename.endswith('.m4a'):
        audio = mutagen.File(filename)
        if '----:com.apple.iTunes:rating wmp' not in audio:
            return None
        rating = int(int(audio['----:com.apple.iTunes:rating wmp'][0].decode()) / 51)
        return rating


def get_unrated(path):
    for file in os.listdir(path):
        if file.split('.')[-1] in ALLOWED_TYPES:
            filepath = os.path.join(path, file)
            if not get_rating(filepath):
                yield filepath


def get_rated(path):
    for file in os.listdir(path):
        if file.split('.')[-1] in ALLOWED_TYPES:
            filepath = os.path.join(path, file)
            rating = get_rating(filepath)
            if rating:
                yield filepath, rating