Add UFID storage style

beets/mediafile.py

@@ -550,22 +550,6 @@ class MP3StorageStyle(StorageStyle):
                     args['lang'] = self.id3_lang
                 mediafile.mgfile.tags.add(mutagen.id3.Frames[self.key](**args))
 
-        # Try to match on "owner" field.
-        elif self.key.startswith('UFID:'):
-            owner = self.key.split(':', 1)[1]
-            frames = mediafile.mgfile.tags.getall(self.key)
-
-            for frame in frames:
-                # Replace existing frame data.
-                if frame.owner == owner:
-                    setattr(frame, self.id3_frame_field, val)
-            else:
-                # New frame.
-                assert isinstance(self.id3_frame_field, str)  # Keyword.
-                frame = mutagen.id3.UFID(owner=owner,
-                    **{self.id3_frame_field: val})
-                mediafile.mgfile.tags.setall('UFID', [frame])
-
         # Just replace based on key.
         else:
             assert isinstance(self.id3_frame_field, str)  # Keyword.
@@ -573,6 +557,24 @@ class MP3StorageStyle(StorageStyle):
                 **{self.id3_frame_field: val})
             mediafile.mgfile.tags.setall(self.key, [frame])
 
+class MP3UFIDStorageStyle(MP3StorageStyle):
+
+    def __init__(self, owner, **kwargs):
+        self.owner = owner
+        super(MP3UFIDStorageStyle, self).__init__('UFID:' + owner, **kwargs)
+
+    def store(self, mediafile, value):
+        frames = mediafile.mgfile.tags.getall(self.key)
+
+        for frame in frames:
+            # Replace existing frame data.
+            if frame.owner == self.owner:
+                setattr(frame, self.id3_frame_field, value)
+        else:
+            # New frame.
+            assert isinstance(self.id3_frame_field, str)  # Keyword.
+            frame = mutagen.id3.UFID(owner=self.owner, data=value)
+            mediafile.mgfile.tags.setall(self.key, [frame])
 
 # The field itself.
 class MediaField(object):
@@ -1198,7 +1200,7 @@ class MediaFile(object):
 
     # MusicBrainz IDs.
     mb_trackid = MediaField(
-        mp3=MP3StorageStyle('UFID:http://musicbrainz.org',
+        mp3=MP3UFIDStorageStyle(owner='http://musicbrainz.org',
                           list_elem=False,
                           id3_frame_field='data'),
         mp4=MP4StorageStyle('----:com.apple.iTunes:MusicBrainz Track Id',

test/test_mediafile.py

@@ -353,14 +353,32 @@ class ReadWriteTest(unittest.TestCase):
             mediafile.original_day = None
             mediafile.save()
 
+            mediafile = beets.mediafile.MediaFile(mediafile.path)
             self.assertEqual(mediafile.rg_track_gain, 0)
             self.assertEqual(mediafile.rg_track_peak, 0)
             self.assertEqual(mediafile.original_year, 0)
             self.assertEqual(mediafile.original_month, 0)
             self.assertEqual(mediafile.original_day, 0)
 
+    def test_read_write_mb_ids(self):
+        for ext in self.extensions:
+            mediafile = full_mediafile_fixture(ext)
+            mediafile.mb_trackid = 'the-id'
+            mediafile.mb_albumid = 'the-id'
+            mediafile.mb_artistid = 'the-id'
+            mediafile.mb_albumartistid = 'the-id'
+            mediafile.mb_releasegroupid = 'the-id'
+            mediafile.save()
 
-def full_mediafile_fixture(ext='mp3'):
+            mediafile = beets.mediafile.MediaFile(mediafile.path)
+            self.assertEqual(mediafile.mb_trackid, 'the-id')
+            self.assertEqual(mediafile.mb_albumid, 'the-id')
+            self.assertEqual(mediafile.mb_artistid, 'the-id')
+            self.assertEqual(mediafile.mb_albumartistid, 'the-id')
+            self.assertEqual(mediafile.mb_releasegroupid, 'the-id')
+
+
+def full_mediafile_fixture(ext):
     """Returns a Mediafile with a lot of tags already set.
     """
     src = os.path.join(_common.RSRC, 'full.{0}'.format(ext))