From e7e7ee64b0d8345182e64f8d1393fbb40344aac5 Mon Sep 17 00:00:00 2001
From: Adrian Sampson <adrian@radbox.org>
Date: Tue, 13 Jul 2010 23:43:49 -0700
Subject: [PATCH] albuminfo() can now take an album id as well as an item

---
 beets/library.py | 24 +++++++++++++++---------
 test/test_db.py  |  5 +++++
 2 files changed, 20 insertions(+), 9 deletions(-)

diff --git a/beets/library.py b/beets/library.py
index 1a72508d2..4cc66e6b0 100644
--- a/beets/library.py
+++ b/beets/library.py
@@ -945,17 +945,23 @@ class Library(BaseLibrary):
     # Album information.
 
     def albuminfo(self, item):
-        # Lazily create a row in the albums table if one doesn't
-        # exist.
-        sql = 'SELECT id FROM albums WHERE artist=? AND album=?'
-        c = self.conn.execute(sql, (item.artist, item.album))
-        row = c.fetchone()
-        if row:
-            album_id = row[0]
+        """Get an album info proxy object given either an item or an
+        album id.
+        """
+        if isinstance(item, int):
+            album_id = item
         else:
-            sql = 'INSERT INTO albums (artist, album) VALUES (?, ?)'
+            # Lazily create a row in the albums table if one doesn't
+            # exist.
+            sql = 'SELECT id FROM albums WHERE artist=? AND album=?'
             c = self.conn.execute(sql, (item.artist, item.album))
-            album_id = c.lastrowid
+            row = c.fetchone()
+            if row:
+                album_id = row[0]
+            else:
+                sql = 'INSERT INTO albums (artist, album) VALUES (?, ?)'
+                c = self.conn.execute(sql, (item.artist, item.album))
+                album_id = c.lastrowid
         return AlbumInfo(self, album_id)
 
     def _album_get(self, album_id, key):
diff --git a/test/test_db.py b/test/test_db.py
index 7f2d72ae7..0d275cf3d 100644
--- a/test/test_db.py
+++ b/test/test_db.py
@@ -409,6 +409,11 @@ class AlbumInfoTest(unittest.TestCase):
         # Cursor should only return one row.
         self.assertNotEqual(c.fetchone(), None)
         self.assertEqual(c.fetchone(), None)
+    
+    def test_albuminfo_by_albumid(self):
+        ai = self.lib.albuminfo(self.i)
+        ai = self.lib.albuminfo(ai.id)
+        self.assertEqual(ai.artist, self.i.artist)
 
 def suite():
     return unittest.TestLoader().loadTestsFromName(__name__)