From 2dc716622ea91f9c6da7626c2643374e7bc1049a Mon Sep 17 00:00:00 2001
From: Kim Wittenburg <5wittenb@informatik.uni-hamburg.de>
Date: Tue, 4 Sep 2018 00:00:33 +0200
Subject: [PATCH] Add Search Prompt

---
 beetsplug/itunes.py | 19 +++++++++++++++++--
 1 file changed, 17 insertions(+), 2 deletions(-)

diff --git a/beetsplug/itunes.py b/beetsplug/itunes.py
index 0897f74..dba8fd2 100644
--- a/beetsplug/itunes.py
+++ b/beetsplug/itunes.py
@@ -7,11 +7,13 @@ import datetime
 import urllib.request
 from tempfile import NamedTemporaryFile
 
+import beets.mediafile as mediafile
 from beets import config
 from beets.autotag.hooks import AlbumInfo, TrackInfo, Distance
+from beets.autotag.match import tag_album
 from beets.plugins import BeetsPlugin, sanitize_choices
 from beets.dbcore import types
-import beets.mediafile as mediafile
+from beets.ui.commands import PromptChoice
 
 AVAILABLE_COUNTRIES = ['DE', 'US', 'GB', 'FR', 'AU', 'CA', 'IT', 'JP', 'DZ', 'AO', 'AI', 'AG', 'AR', 'AM', 'AT', 'AZ', 'BS', 'BH', 'BD', 'BB', 'BY', 'BE', 'BZ', 'BM', 'BO', 'BW', 'BR', 'BN', 'BG', 'CM', 'KY', 'CL', 'CN', 'CO', 'CR', 'CI', 'HR', 'CY', 'CZ', 'DK', 'DM', 'DO', 'EC', 'EG', 'SV', 'EE', 'ET', 'FI', 'GH', 'GR', 'GD', 'GT', 'GY', 'HN', 'HK', 'HU', 'IS', 'IN', 'ID', 'IE', 'IL', 'JM', 'JO', 'KZ', 'KE', 'KR', 'KW', 'LV', 'LB', 'LY', 'LI', 'LT', 'LU', 'MO', 'MK', 'MG', 'MY', 'MV', 'ML', 'MT', 'MU', 'MX', 'MD', 'MS', 'MM', 'NP', 'NL', 'NZ', 'NI', 'NE', 'NG', 'NO', 'OM', 'PK', 'PA', 'PY', 'PE', 'PH', 'PL', 'PT', 'QA', 'RO', 'RU', 'KN', 'LC', 'VC', 'SA', 'SN', 'RS', 'SG', 'SK', 'SI', 'ZA', 'ES', 'LK', 'SR', 'SE', 'CH', 'TW', 'TZ', 'TH', 'TT', 'TN', 'TR', 'TC', 'UG', 'UA', 'AE', 'UY', 'UZ', 'VE', 'VN', 'VG', 'YE']
 VA_ARTIST_ID = 120644327
@@ -73,8 +75,11 @@ class iTunesPlugin(BeetsPlugin):
             'artwork': False, # Whether to also download artwork from iTunes. Does not work well with fetchart
             'method': 'lookup', # Valid values: 'search', 'lookup', 'both', 'lookup (search fallback)'
         })
+        self.force_search = False
         if self.method in ['lookup', 'lookup (search fallback)', 'both']:
             self.register_listener('import_task_start', self.scan_track_ids)
+        if self.method in ['lookup', 'lookup (search fallback)']:
+            self.register_listener('before_choose_candidate', self.add_prompt_choices)
         self.register_listener('import_task_apply', self.apply_itunes_metadata)
         if self.use_itunes_artwork:
             self.import_stages += [ self.fetch_artwork ]
@@ -119,6 +124,8 @@ class iTunesPlugin(BeetsPlugin):
         """Returns a list of AlbumInfo objects for iTunes search results
         matching an album and artist.
         """
+        force_search = self.force_search
+        self.force_search = False
         try:
             candidates = []
             # Lookup the track ids assigned by iTunes Match
@@ -130,7 +137,7 @@ class iTunesPlugin(BeetsPlugin):
                 for item in items:
                     _track_ids.pop(item, None)
 
-            if self.method in ['search', 'both'] or (self.method == 'lookup (search fallback)' and len(candidates) == 0):
+            if force_search or self.method in ['search', 'both'] or (self.method == 'lookup (search fallback)' and len(candidates) == 0):
                 candidates += self.search_albums(artist, album, va_likely)
             return candidates
         except ConnectionError:
@@ -319,6 +326,14 @@ class iTunesPlugin(BeetsPlugin):
         info.itunes_advisory = decode_explicitness(track.track_explicitness)
         info.itunes_artwork_url = decode_artwork_url(track.artwork_url_100)
         return info
+    
+    def add_prompt_choices(self, session, task):
+        return [PromptChoice('c', 'searCh iTunes', self.prompt_search)]
+    
+    def prompt_search(self, session, task):
+        self.force_search = True
+        _, _, proposal = tag_album(task.items)
+        return proposal
 
 def decode_explicitness(advisory_string):
     return ['notExplicit', 'explicit', 'cleaned'].index(advisory_string)