diff --git a/beetsplug/lastgenre/__init__.py b/beetsplug/lastgenre/__init__.py index dcf5b2fe2..581772519 100644 --- a/beetsplug/lastgenre/__init__.py +++ b/beetsplug/lastgenre/__init__.py @@ -38,7 +38,6 @@ from beets import library log = logging.getLogger('beets') LASTFM = pylast.LastFMNetwork(api_key=plugins.LASTFM_KEY) -C14N_TREE = os.path.join(os.path.dirname(__file__), 'genres-tree.yaml') PYLAST_EXCEPTIONS = ( pylast.WSError, @@ -113,12 +112,16 @@ def find_parents(candidate, branches): # Main plugin logic. +WHITELIST = os.path.join(os.path.dirname(__file__), 'genres.txt') +C14N_TREE = os.path.join(os.path.dirname(__file__), 'genres-tree.yaml') + + class LastGenrePlugin(plugins.BeetsPlugin): def __init__(self): super(LastGenrePlugin, self).__init__() self.config.add({ - 'whitelist': os.path.join(os.path.dirname(__file__), 'genres.txt'), + 'whitelist': None, 'min_weight': 10, 'count': 1, 'fallback': None, @@ -139,13 +142,20 @@ class LastGenrePlugin(plugins.BeetsPlugin): self._genre_cache = {} - # Read the whitelist file. + # Read the whitelist file if enabled. self.whitelist = set() - with open(self.config['whitelist'].as_filename()) as f: - for line in f: - line = line.decode('utf8').strip().lower() - if line and not line.startswith(u'#'): - self.whitelist.add(line) + wl_filename = self.config['whitelist'].get() + if wl_filename is not None: + wl_filename = wl_filename.strip() + if not wl_filename: + wl_filename = WHITELIST + wl_filename = normpath(wl_filename) + with open(wl_filename, 'r') as f: + # with open(self.config['whitelist'].as_filename()) as f: + for line in f: + line = line.decode('utf8').strip().lower() + if line and not line.startswith(u'#'): + self.whitelist.add(line) # Read the genres tree for canonicalization if enabled. self.c14n_branches = [] @@ -184,13 +194,18 @@ class LastGenrePlugin(plugins.BeetsPlugin): # Extend the list to consider tags parents in the c14n tree tags_all = [] for tag in tags: - parents = [x for x in find_parents(tag, self.c14n_branches) if - self._is_allowed(x)] + # Add parents that are in the whitelist, or add the oldest + # ancestor if no whitelist + if self.whitelist: + parents = [x for x in find_parents(tag, self.c14n_branches) + if self._is_allowed(x)] + else: + parents = [find_parents(tag, self.c14n_branches)[-1]] + tags_all += parents if len(tags_all) >= count: break tags = tags_all - # c14n only adds allowed genres but we may have had forbidden genres in # the original tags list tags = [x.title() for x in tags if self._is_allowed(x)] diff --git a/test/test_lastgenre.py b/test/test_lastgenre.py index b4fa9cc4d..89fb6eec9 100644 --- a/test/test_lastgenre.py +++ b/test/test_lastgenre.py @@ -27,30 +27,48 @@ lastGenrePlugin = lastgenre.LastGenrePlugin() class LastGenrePluginTest(unittest.TestCase): - def _setup_config(self, whitelist=set(), canonical=None, count=1): + def _setup_config(self, whitelist=None, canonical=None, count=1): config['lastgenre']['canonical'] = canonical config['lastgenre']['count'] = count - config['lastgenre']['whitelist'] = \ - os.path.join(os.path.dirname(lastgenre.__file__), 'genres.txt') + if whitelist == '': + config['lastgenre']['whitelist'] = whitelist lastGenrePlugin.setup() if whitelist: lastGenrePlugin.whitelist = whitelist - def test_defaults(self): - """Tests whitelist and c14n options with default filepaths + def test_default(self): + """Fetch genres with whitelist and c14n deactivated """ - # default whitelist and c14n - self._setup_config(canonical='') - self.assertEqual(lastGenrePlugin._resolve_genres(['delta blues']), - 'Blues') - self.assertEqual(lastGenrePlugin._resolve_genres(['iota blues']), '') - - # default whitelist and no c14n self._setup_config() self.assertEqual(lastGenrePlugin._resolve_genres(['delta blues']), - 'delta blues') + 'Delta Blues') - def test_whitelist(self): + def test_c14n_only(self): + """Default c14n tree funnels up to most common genre except for *wrong* + genres that stay unchanged. + """ + self._setup_config(canonical='', count=99) + self.assertEqual(lastGenrePlugin._resolve_genres(['delta blues']), + 'Blues') + self.assertEqual(lastGenrePlugin._resolve_genres(['iota blues']), + 'Iota Blues') + + def test_whitelist_only(self): + """Default whitelist rejects *wrong* (non existing) genres. + """ + self._setup_config(whitelist='') + self.assertEqual(lastGenrePlugin._resolve_genres(['iota blues']), + '') + + def test_whitelist_c14n(self): + """Default whitelist and c14n both activated result in all parents + genres being selected (from specific to common). + """ + self._setup_config(canonical='', whitelist='', count=99) + self.assertEqual(lastGenrePlugin._resolve_genres(['delta blues']), + 'Delta Blues, Country Blues, Blues') + + def test_whitelist_custom(self): """Keep only genres that are in the whitelist. """ self._setup_config(whitelist=set(['blues', 'rock', 'jazz']), @@ -78,13 +96,19 @@ class LastGenrePluginTest(unittest.TestCase): self._setup_config(whitelist=set(['blues', 'rock', 'jazz']), canonical='', count=2) - # thanks to c14n, 'blues' superseeds 'country blues' and takes the # second slot self.assertEqual(lastGenrePlugin._resolve_genres( ['jazz', 'pop', 'country blues', 'rock']), 'Jazz, Blues') + def test_c14n_whitelist(self): + """Genres first pass through c14n and are then filtered + """ + self._setup_config(canonical='', whitelist=set(['rock'])) + self.assertEqual(lastGenrePlugin._resolve_genres(['delta blues']), + '') + def suite(): return unittest.TestLoader().loadTestsFromName(__name__)