symuid/symuid/tag_interface.py

# -*- coding: utf-8 -*-

import mutagen.id3
import mutagen.mp4
import re


class _mutagen:

    @property
    def track_path(self):
        return self._mutagen_file.filename

    def save(self):
        self._mutagen_file.save()


class ID3(_mutagen):

    def __init__(self, mutagen_file):
        assert mutagen_file.tags, mutagen_file
        assert isinstance(mutagen_file.tags, mutagen.id3.ID3), \
            mutagen_file.tags
        self._mutagen_file = mutagen_file

    def get_free_ints(self, tag_label_prefix):
        for t in self._mutagen_file.tags.getall('TXXX:' + tag_label_prefix):
            assert len(t.text) == 1, t
            yield (t.desc, int(t.text[0]))

    def get_free_int(self, tag_label):
        tags = [t for t in self.get_free_ints(tag_label)]
        if len(tags) == 0:
            raise KeyError(tag_label)
        else:
            assert len(tags) == 1, tags
            assert tags[0][0] == tag_label, tag
            return tags[0][1]

    def set_free_int(self, tag_label, data):
        # mutagen.id3._specs.EncodedTextSpec.write encodes 'desc' and 'text'
        tag = mutagen.id3.TXXX(
            encoding=mutagen.id3.Encoding.LATIN1,
            desc=tag_label,
            text=[str(data)],
        )
        self._mutagen_file.tags.add(tag)
        return tag


class MP4(_mutagen):

    def __init__(self, mutagen_file):
        assert mutagen_file.tags, mutagen_file
        assert isinstance(mutagen_file.tags, mutagen.mp4.MP4Tags), \
            mutagen_file.tags
        self._mutagen_file = mutagen_file

    @staticmethod
    def _freeform_to_int(freeform):
        # "a signed big-endian integer with length one of { 1,2,3,4,8 } bytes"
        assert freeform.dataformat == mutagen.mp4.AtomDataType.INTEGER, freeform
        return int.from_bytes(freeform, byteorder='big', signed=True)

    def get_free_ints(self, tag_label_prefix):
        label_pattern = re.compile(r'^----:{}(:|$)'.format(
            re.escape(tag_label_prefix),
        ))
        for label, values in self._mutagen_file.tags.items():
            if label_pattern.match(label):
                assert len(values) == 1, (label, values)
                value = MP4._freeform_to_int(values[0])
                yield (re.sub(r'^----:', '', label), value)

    def get_free_int(self, tag_label):
        tag, = self._mutagen_file.tags['----:' + tag_label]
        return MP4._freeform_to_int(tag)

    def set_free_int(self, tag_label, data):
        assert isinstance(data, int)
        tag = mutagen.mp4.MP4FreeForm(
            # "a signed big-endian integer with length one of { 1,2,3,4,8 } bytes"
            dataformat=mutagen.mp4.AtomDataType.INTEGER,
            # TODO set byte length properly
            data=data.to_bytes(1, byteorder='big', signed=True),
        )
        self._mutagen_file.tags['----:' + tag_label] = tag
        return tag