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

    # http://id3.org/id3v2.4.0-frames#4.1.
    _UFID_OWNER_ID = 'symuid'

    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)],
        )
        # TODO overwrite instead of add() ?
        self._mutagen_file.tags.add(tag)
        return tag

    def get_track_uuid(self):
        for ufid in self._mutagen_file.tags.getall('UFID'):
            if ufid.owner == self._UFID_OWNER_ID:
                return ufid.data
        return None

    def set_track_uuid(self, uuid):
        # mutagen.id3._specs.EncodedTextSpec.write encodes 'owner'
        tag = mutagen.id3.UFID(owner=self._UFID_OWNER_ID, data=uuid)
        self._mutagen_file.tags.add(tag)
        return tag


class MP4(_mutagen):

    _UUID_TAG_KEY = 'symuid:uuid'

    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():
            # TODO overwrite instead of add() ?
            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(self, tag_label):
        # freeform keys start with '----'
        # http://mutagen.readthedocs.io/en/latest/api/mp4.html
        tags = self._mutagen_file.tags['----:' + tag_label]
        assert len(tags) == 1, tags
        return tags[0]

    def get_free_int(self, tag_label):
        return MP4._freeform_to_int(self._get_free(tag_label))

    def _get_free_uuid(self, tag_label):
        tag = self._get_free(tag_label)
        assert tag.dataformat == mutagen.mp4.AtomDataType.UUID, tag.dataformat
        return 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

    def _set_free_uuid(self, tag_label, data):
        assert isinstance(data, bytes)
        tag = mutagen.mp4.MP4FreeForm(
            data=data,
            # https://mutagen.readthedocs.io/en/latest/api/mp4.html#mutagen.mp4.AtomDataType.UUID
            dataformat=mutagen.mp4.AtomDataType.UUID,
        )
        self._mutagen_file.tags['----:' + tag_label] = [tag]
        return tag

    def get_track_uuid(self):
        try:
            return self._get_free_uuid(self._UUID_TAG_KEY)
        except KeyError:
            return None

    def set_track_uuid(self, uuid):
        return self._set_free_uuid(self._UUID_TAG_KEY, uuid)