multi-disc album collapsing based on heuristics (#42)

beets/autotag/__init__.py

@@ -16,9 +16,10 @@
 """
 import os
 import logging
+import re
 
 from beets import library, mediafile
-from beets.util import sorted_walk
+from beets.util import sorted_walk, ancestry
 
 # Parts of external interface.
 from .hooks import AlbumInfo, TrackInfo
@@ -30,6 +31,10 @@ from .match import STRONG_REC_THRESH, MEDIUM_REC_THRESH, REC_GAP_THRESH
 # Global logger.
 log = logging.getLogger('beets')
 
+# Constants for directory walker.
+MULTIDISC_MARKERS = (r'part', r'volume', r'vol\.', r'disc', r'cd')
+MULTIDISC_PAT_FMT = r'%s\s*\d'
+
 
 # Additional utilities for the main interface.
 
@@ -40,6 +45,9 @@ def albums_in_dir(path, ignore=()):
     containing any media files is an album. Directories and file names
     that match the glob patterns in ``ignore`` are skipped.
     """
+    collapse_root = None
+    collapse_items = None
+
     for root, dirs, files in sorted_walk(path, ignore):
         # Get a list of items in the directory.
         items = []
@@ -52,11 +60,47 @@ def albums_in_dir(path, ignore=()):
                 log.warn('unreadable file: ' + filename)
             else:
                 items.append(i)
+
+        # If we're collapsing, test to see whether we should continue to
+        # collapse. If so, just add to the collapsed item set;
+        # otherwise, end the collapse and continue as normal.
+        if collapse_root is not None:
+            if collapse_root in ancestry(root):
+                # Still collapsing.
+                collapse_items += items
+                continue
+            else:
+                # Collapse finished. Yield the collapsed directory and
+                # proceed to process the current one.
+                yield collapse_root, collapse_items
+                collapse_root = collapse_items = None
+
+        # Does the current directory look like a multi-disc album? If
+        # so, begin collapsing here.
+        if dirs and not items: # Must be only directories.
+            multidisc = False
+            for marker in MULTIDISC_MARKERS:
+                pat = MULTIDISC_PAT_FMT % marker
+                if all(re.search(pat, dirname, re.I) for dirname in dirs):
+                    multidisc = True
+                    break
+
+            # This becomes True only when all directories match a
+            # pattern for a single marker.
+            if multidisc:
+                # Start collapsing; continue to the next iteration.
+                collapse_root = root
+                collapse_items = []
+                continue
         
         # If it's nonempty, yield it.
         if items:
             yield root, items
 
+    # Clear out any unfinished collapse.
+    if collapse_root is not None:
+        yield collapse_root, collapse_items
+
 def apply_item_metadata(item, track_info):
     """Set an item's metadata from its matched TrackInfo object.
     """

test/test_autotag.py

@@ -197,6 +197,43 @@ class AlbumsInDirTest(unittest.TestCase):
             else:
                 self.assertEqual(len(album), 1)
 
+class MultiDiscAlbumsInDirTest(unittest.TestCase):
+    def setUp(self):
+        self.base = os.path.abspath(os.path.join(_common.RSRC, 'tempdir'))
+        os.mkdir(self.base)
+        
+        os.mkdir(os.path.join(self.base, 'album1'))
+        os.mkdir(os.path.join(self.base, 'album1', 'disc 1'))
+        os.mkdir(os.path.join(self.base, 'album1', 'disc 2'))
+
+        os.mkdir(os.path.join(self.base, 'dir2'))
+        os.mkdir(os.path.join(self.base, 'dir2', 'disc 1'))
+        os.mkdir(os.path.join(self.base, 'dir2', 'something'))
+        
+        _mkmp3(os.path.join(self.base, 'album1', 'disc 1', 'song1.mp3'))
+        _mkmp3(os.path.join(self.base, 'album1', 'disc 2', 'song2.mp3'))
+        _mkmp3(os.path.join(self.base, 'album1', 'disc 2', 'song3.mp3'))
+
+        _mkmp3(os.path.join(self.base, 'dir2', 'disc 1', 'song4.mp3'))
+        _mkmp3(os.path.join(self.base, 'dir2', 'something', 'song5.mp3'))
+
+    def tearDown(self):
+        shutil.rmtree(self.base)
+
+    def test_coalesce_multi_disc_album(self):
+        albums = list(autotag.albums_in_dir(self.base))
+        self.assertEquals(len(albums), 3)
+        root, items = albums[0]
+        self.assertEquals(root, os.path.join(self.base, 'album1'))
+        self.assertEquals(len(items), 3)
+
+    def test_separate_red_herring(self):
+        albums = list(autotag.albums_in_dir(self.base))
+        root, items = albums[1]
+        self.assertEquals(root, os.path.join(self.base, 'dir2', 'disc 1'))
+        root, items = albums[2]
+        self.assertEquals(root, os.path.join(self.base, 'dir2', 'something'))
+
 class OrderingTest(unittest.TestCase):
     def item(self, title, track):
         return Item({