Add musicbrainz id option to importer

beets/autotag/match.py

@@ -370,7 +370,7 @@ def _add_candidate(items, results, info):
 
 
 def tag_album(items, search_artist=None, search_album=None,
-              search_id=None):
+              search_ids=[]):
     """Return a tuple of a artist name, an album name, a list of
     `AlbumMatch` candidates from the metadata backend, and a
     `Recommendation`.
@@ -397,9 +397,11 @@ def tag_album(items, search_artist=None, search_album=None,
     candidates = {}
 
     # Search by explicit ID.
-    if search_id is not None:
-        log.debug(u'Searching for album ID: {0}', search_id)
-        search_cands = hooks.albums_for_id(search_id)
+    if search_ids:
+        search_cands = []
+        for search_id in search_ids:
+            log.debug(u'Searching for album ID: {0}', search_id)
+            search_cands.extend(hooks.albums_for_id(search_id))
 
     # Use existing metadata or text search.
     else:
@@ -444,33 +446,35 @@ def tag_album(items, search_artist=None, search_album=None,
 
 
 def tag_item(item, search_artist=None, search_title=None,
-             search_id=None):
+             search_ids=[]):
     """Attempts to find metadata for a single track. Returns a
     `(candidates, recommendation)` pair where `candidates` is a list of
     TrackMatch objects. `search_artist` and `search_title` may be used
     to override the current metadata for the purposes of the MusicBrainz
-    title; likewise `search_id`.
+    title; likewise `search_ids`.
     """
     # Holds candidates found so far: keys are MBIDs; values are
     # (distance, TrackInfo) pairs.
     candidates = {}
 
     # First, try matching by MusicBrainz ID.
-    trackid = search_id or item.mb_trackid
-    if trackid:
-        log.debug(u'Searching for track ID: {0}', trackid)
-        for track_info in hooks.tracks_for_id(trackid):
-            dist = track_distance(item, track_info, incl_artist=True)
-            candidates[track_info.track_id] = \
-                hooks.TrackMatch(dist, track_info)
-            # If this is a good match, then don't keep searching.
-            rec = _recommendation(candidates.values())
-            if rec == Recommendation.strong and not config['import']['timid']:
-                log.debug(u'Track ID match.')
-                return candidates.values(), rec
+    trackids = search_ids or filter(None, [item.mb_trackid])
+    if trackids:
+        for trackid in trackids:
+            log.debug(u'Searching for track ID: {0}', trackid)
+            for track_info in hooks.tracks_for_id(trackid):
+                dist = track_distance(item, track_info, incl_artist=True)
+                candidates[track_info.track_id] = \
+                    hooks.TrackMatch(dist, track_info)
+                # If this is a good match, then don't keep searching.
+                rec = _recommendation(candidates.values())
+                if rec == Recommendation.strong and \
+                        not config['import']['timid']:
+                    log.debug(u'Track ID match.')
+                    return candidates.values(), rec
 
     # If we're searching by ID, don't proceed.
-    if search_id is not None:
+    if search_ids != []:
         if candidates:
             return candidates.values(), rec
         else:

beets/importer.py

@@ -581,8 +581,13 @@ def handle_created(self, session):
     def lookup_candidates(self):
         """Retrieve and store candidates for this album.
         """
+        # Use a MusicBrainz id directly if provided by the importer -m option.
+        mb_ids = []
+        if config['import']['musicbrainz_ids'].exists():
+            mb_ids = config['import']['musicbrainz_ids'].get()
+
         artist, album, candidates, recommendation = \
-            autotag.tag_album(self.items)
+            autotag.tag_album(self.items, search_ids=mb_ids)
         self.cur_artist = artist
         self.cur_album = album
         self.candidates = candidates
@@ -821,7 +826,13 @@ def _emit_imported(self, lib):
             plugins.send('item_imported', lib=lib, item=item)
 
     def lookup_candidates(self):
-        candidates, recommendation = autotag.tag_item(self.item)
+        # Use a MusicBrainz id directly if provided by the importer -m option.
+        mb_ids = []
+        if config['import']['musicbrainz_ids'].exists():
+            mb_ids = config['import']['musicbrainz_ids'].get()
+
+        candidates, recommendation = autotag.tag_item(self.item,
+                                                      search_ids=mb_ids)
         self.candidates = candidates
         self.rec = recommendation
 

beets/ui/commands.py

@@ -734,7 +734,7 @@ def choose_match(self, task):
                 search_id = manual_id(False)
                 if search_id:
                     _, _, candidates, rec = autotag.tag_album(
-                        task.items, search_id=search_id
+                        task.items, search_ids=[search_id]
                     )
             elif choice in extra_ops.keys():
                 # Allow extra ops to automatically set the post-choice.
@@ -787,7 +787,7 @@ def choose_item(self, task):
                 search_id = manual_id(True)
                 if search_id:
                     candidates, rec = autotag.tag_item(task.item,
-                                                       search_id=search_id)
+                                                       search_ids=[search_id])
             elif choice in extra_ops.keys():
                 # Allow extra ops to automatically set the post-choice.
                 post_choice = extra_ops[choice](self, task)
@@ -1022,6 +1022,10 @@ def import_func(lib, opts, args):
     '--pretend', dest='pretend', action='store_true',
     help='just print the files to import'
 )
+import_cmd.parser.add_option(
+    '-m', '--musicbrainzid', dest='musicbrainz_ids', action='append',
+    help='restrict the matching to a single MusicBrainz id'
+)
 import_cmd.func = import_func
 default_commands.append(import_cmd)
 

beetsplug/bench.py

@@ -75,7 +75,7 @@ def match_benchmark(lib, prof, query=None, album_id=None):
 
     # Run the match.
     def _run_match():
-        match.tag_album(items, search_id=album_id)
+        match.tag_album(items, search_ids=[album_id])
     if prof:
         cProfile.runctx('_run_match()', {}, {'_run_match': _run_match},
                         'match.prof')

test/test_importer.py

@@ -1689,6 +1689,84 @@ def test_import_pretend_empty(self):
                          .format(displayable_path(self.empty_path))])
 
 
+class ImportMusicBrainzIdTest(_common.TestCase, ImportHelper):
+    """Test the --musicbrainzid argument."""
+    # The Beatles - Revolver (14 tracks)
+    ID_REVOLVER = ('https://musicbrainz.org/release/'
+                   'd74365b0-288d-3b85-b344-7f9f603378b6')
+    # The Beatles - Rubber Soul (14 tracks)
+    ID_RUBBER = ('https://musicbrainz.org/release/'
+                 'd5d19748-72fc-3f20-aae7-7ca22da283a0')
+    # The Beatles - Rubber Soul - Michelle
+    ID_MICHELLE = ('https://musicbrainz.org/recording/'
+                   '45846566-fb79-4758-a006-15cf69e0d138')
+    # The Beatles - Rubber Soul - Girl
+    ID_GIRL = ('https://musicbrainz.org/recording/'
+               '91ec42bc-a04d-497e-b829-439831c8dbcb')
+
+    def setUp(self):
+        self.setup_beets()
+        self._create_import_dir(1)
+
+    def tearDown(self):
+        self.teardown_beets()
+
+    def test_one_mbid_one_album(self):
+        self.config['import']['musicbrainz_ids'] = [self.ID_REVOLVER]
+        self._setup_import_session()
+        self.importer.add_choice(importer.action.APPLY)
+        self.importer.run()
+        self.assertEqual(self.lib.albums().get().album, 'Revolver')
+
+    def test_several_mbid_one_album(self):
+        self.config['import']['musicbrainz_ids'] = [self.ID_REVOLVER,
+                                                    self.ID_RUBBER]
+        self._setup_import_session()
+        self.importer.add_choice(2)
+        self.importer.add_choice(importer.action.APPLY)
+        self.importer.run()
+        self.assertEqual(self.lib.albums().get().album, 'Rubber Soul')
+
+    def test_one_mbid_one_singleton(self):
+        self.config['import']['musicbrainz_ids'] = [self.ID_MICHELLE]
+        self._setup_import_session(singletons=True)
+        self.importer.add_choice(importer.action.APPLY)
+        self.importer.run()
+        self.assertEqual(self.lib.items().get().title, 'Michelle')
+
+    def test_several_mbid_one_singleton(self):
+        self.config['import']['musicbrainz_ids'] = [self.ID_MICHELLE,
+                                                    self.ID_GIRL]
+        self._setup_import_session(singletons=True)
+        self.importer.add_choice(2)
+        self.importer.add_choice(importer.action.APPLY)
+        self.importer.run()
+        self.assertEqual(self.lib.items().get().title, 'Girl')
+
+    def test_candidates_album(self):
+        """Test directly task.lookup_candidates()."""
+        self.config['import']['musicbrainz_ids'] = [self.ID_REVOLVER,
+                                                    self.ID_RUBBER,
+                                                    'invalid id']  # discarded
+        task = importer.ImportTask(paths=self.import_dir,
+                                   toppath='top path',
+                                   items=[_common.item()])
+        task.lookup_candidates()
+        self.assertEqual(set(['Revolver', 'Rubber Soul']),
+                         set([c.info.album for c in task.candidates]))
+
+    def test_candidates_singleton(self):
+        """Test directly task.lookup_candidates()."""
+        self.config['import']['musicbrainz_ids'] = [self.ID_MICHELLE,
+                                                    self.ID_GIRL,
+                                                    'invalid id']  # discarded
+        task = importer.SingletonImportTask(toppath='top path',
+                                            item=_common.item())
+        task.lookup_candidates()
+        self.assertEqual(set(['Michelle', 'Girl']),
+                         set([c.info.title for c in task.candidates]))
+
+
 def suite():
     return unittest.TestLoader().loadTestsFromName(__name__)