Adapting FolderManager and the folder part of the CLI

cli.py

@@ -22,6 +22,11 @@
 import sys, cmd, argparse, getpass, time
 import config
 
+from db import get_store, Folder
+from managers.folder import FolderManager
+from managers.user import UserManager
+from scanner import Scanner
+
 class CLIParser(argparse.ArgumentParser):
 	def error(self, message):
 		self.print_usage(sys.stderr)
@@ -53,7 +58,7 @@ class CLI(cmd.Cmd):
 
 		return method
 
-	def __init__(self):
+	def __init__(self, store):
 		cmd.Cmd.__init__(self)
 
 		# Generate do_* and help_* methods
@@ -69,6 +74,8 @@ class CLI(cmd.Cmd):
 				for action, subparser in getattr(self.__class__, command + '_subparsers').choices.iteritems():
 					setattr(self, 'help_{} {}'.format(command, action), subparser.print_help)
 
+		self.__store = store
+
 	def do_EOF(self, line):
 		return True
 
@@ -105,17 +112,17 @@ class CLI(cmd.Cmd):
 
 	def folder_list(self):
 		print 'Name\t\tPath\n----\t\t----'
-		print '\n'.join('{0: <16}{1}'.format(f.name, f.path) for f in db.Folder.query.filter(db.Folder.root == True))
+		print '\n'.join('{0: <16}{1}'.format(f.name, f.path) for f in self.__store.find(Folder, Folder.root == True))
 
 	def folder_add(self, name, path):
-		ret = FolderManager.add(name, path)
+		ret = FolderManager.add(self.__store, name, path)
 		if ret != FolderManager.SUCCESS:
 			print FolderManager.error_str(ret)
 		else:
 			print "Folder '{}' added".format(name)
 
 	def folder_delete(self, name):
-		ret = FolderManager.delete_by_name(name)
+		ret = FolderManager.delete_by_name(self.__store, name)
 		if ret != FolderManager.SUCCESS:
 			print FolderManager.error_str(ret)
 		else:
@@ -134,19 +141,19 @@ class CLI(cmd.Cmd):
 					print "Scanning '{0}': {1}% ({2}/{3})".format(self.__name, (scanned * 100) / total, scanned, total)
 					self.__last_display = time.time()
 
-		s = Scanner(db.session)
+		s = Scanner(self.__store)
 		if folders:
-			folders = map(lambda n: db.Folder.query.filter(db.Folder.name == n and db.Folder.root == True).first() or n, folders)
+			folders = map(lambda n: self.__store.find(Folder, Folder.name == n, Folder.root == True).one() or n, folders)
 			if any(map(lambda f: isinstance(f, basestring), folders)):
 				print "No such folder(s): " + ' '.join(f for f in folders if isinstance(f, basestring))
-			for folder in filter(lambda f: isinstance(f, db.Folder), folders):
-				FolderManager.scan(folder.id, s, TimedProgressDisplay(folder.name))
+			for folder in filter(lambda f: isinstance(f, Folder), folders):
+				FolderManager.scan(self.__store, folder.id, s, TimedProgressDisplay(folder.name))
 		else:
-			for folder in db.Folder.query.filter(db.Folder.root == True):
-				FolderManager.scan(folder.id, s, TimedProgressDisplay(folder.name))
+			for folder in self.__store.find(Folder, Folder.root == True):
+				FolderManager.scan(self.__store, folder.id, s, TimedProgressDisplay(folder.name))
 
 		added, deleted = s.stats()
-		db.session.commit()
+		self.__store.commit()
 
 		print "Scanning done"
 		print 'Added: %i artists, %i albums, %i tracks' % (added[0], added[1], added[2])
@@ -219,15 +226,9 @@ if __name__ == "__main__":
 	if not config.check():
 		sys.exit(1)
 
-	import db
-	db.init_db()
-
-	from managers.folder import FolderManager
-	from managers.user import UserManager
-	from scanner import Scanner
-
+	cli = CLI(get_store(config.get('base', 'database_uri')))
 	if len(sys.argv) > 1:
-		CLI().onecmd(' '.join(sys.argv[1:]))
+		cli.onecmd(' '.join(sys.argv[1:]))
 	else:
-		CLI().cmdloop()
+		cli.cmdloop()
 

db.py

@@ -19,7 +19,10 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 from storm.properties import *
-from storm.references import *
+from storm.references import Reference, ReferenceSet
+from storm.database import create_database
+from storm.store import Store
+from storm.variables import Variable
 
 import uuid, datetime, time
 import os.path
@@ -27,6 +30,18 @@ import os.path
 def now():
 	return datetime.datetime.now().replace(microsecond = 0)
 
+class UnicodeOrStrVariable(Variable):
+	__slots__ = ()
+
+	def parse_set(self, value, from_db):
+		if isinstance(value, unicode):
+			return value
+		elif isinstance(value, str):
+			return unicode(value)
+		raise TypeError("Expected unicode, found %r: %r" % (type(value), value))
+
+Unicode.variable_class = UnicodeOrStrVariable
+
 class Folder(object):
 	__storm_table__ = 'folder'
 
@@ -353,3 +368,8 @@ class PlaylistTrack(object):
 
 Playlist.tracks = ReferenceSet(Playlist.id, PlaylistTrack.playlist_id, PlaylistTrack.track_id, Track.id)
 
+def get_store(database_uri):
+	database = create_database(database_uri)
+	store = Store(database)
+	return store
+

managers/folder.py

@@ -19,7 +19,7 @@
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
 import os.path, uuid
-from db import Folder, Artist, session
+from db import Folder, Artist, Album, Track
 
 class FolderManager:
 	SUCCESS = 0
@@ -30,7 +30,7 @@ class FolderManager:
 	NO_SUCH_FOLDER = 5
 
 	@staticmethod
-	def get(uid):
+	def get(store, uid):
 		if isinstance(uid, basestring):
 			try:
 				uid = uuid.UUID(uid)
@@ -41,33 +41,36 @@ class FolderManager:
 		else:
 			return FolderManager.INVALID_ID, None
 
-		folder = Folder.query.get(uid)
+		folder = store.get(Folder, uid)
 		if not folder:
 			return FolderManager.NO_SUCH_FOLDER, None
 
 		return FolderManager.SUCCESS, folder
 
 	@staticmethod
-	def add(name, path):
-		if Folder.query.filter(Folder.name == name and Folder.root == True).first():
+	def add(store, name, path):
+		if not store.find(Folder, Folder.name == name, Folder.root == True).is_empty():
 			return FolderManager.NAME_EXISTS
 
 		path = os.path.abspath(path)
 		if not os.path.isdir(path):
 			return FolderManager.INVALID_PATH
-		folder = Folder.query.filter(Folder.path == path).first()
-		if folder:
+		if not store.find(Folder, Folder.path == path).is_empty():
 			return FolderManager.PATH_EXISTS
 
-		folder = Folder(root = True, name = name, path = path)
-		session.add(folder)
-		session.commit()
+		folder = Folder()
+		folder.root = True
+		folder.name = name
+		folder.path = path
+
+		store.add(folder)
+		store.commit()
 
 		return FolderManager.SUCCESS
 
 	@staticmethod
-	def delete(uid):
-		status, folder = FolderManager.get(uid)
+	def delete(store, uid):
+		status, folder = FolderManager.get(store, uid)
 		if status != FolderManager.SUCCESS:
 			return status
 
@@ -75,37 +78,37 @@ class FolderManager:
 			return FolderManager.NO_SUCH_FOLDER
 
 		# delete associated tracks and prune empty albums/artists
-		for artist in Artist.query.all():
-			for album in artist.albums[:]:
-				for track in filter(lambda t: t.root_folder.id == folder.id, album.tracks):
-					album.tracks.remove(track)
-					session.delete(track)
-				if len(album.tracks) == 0:
-					artist.albums.remove(album)
-					session.delete(album)
-			if len(artist.albums) == 0:
-				session.delete(artist)
+		potentially_removed_albums = set()
+		for track in store.find(Track, Track.root_folder_id == folder.id):
+			potentially_removed_albums.add(track.album)
+			store.remove(track)
+		potentially_removed_artists = set()
+		for album in filter(lambda album: album.tracks.count() == 0, potentially_removed_albums):
+			potentially_removed_artists.add(album.artist)
+			store.remove(album)
+		for artist in filter(lambda artist: artist.albums.count() == 0, potentially_removed_artists):
+			store.remove(artist)
 
 		def cleanup_folder(folder):
 			for f in folder.children:
 				cleanup_folder(f)
-			session.delete(folder)
+			store.remove(folder)
 
 		cleanup_folder(folder)
-		session.commit()
+		store.commit()
 
 		return FolderManager.SUCCESS
 
 	@staticmethod
-	def delete_by_name(name):
-		folder = Folder.query.filter(Folder.name == name and Folder.root == True).first()
+	def delete_by_name(store, name):
+		folder = store.find(Folder, Folder.name == name, Folder.root == True).one()
 		if not folder:
 			return FolderManager.NO_SUCH_FOLDER
-		return FolderManager.delete(folder.id)
+		return FolderManager.delete(store, folder.id)
 
 	@staticmethod
-	def scan(uid, scanner, progress_callback = None):
-		status, folder = FolderManager.get(uid)
+	def scan(store, uid, scanner, progress_callback = None):
+		status, folder = FolderManager.get(store, uid)
 		if status != FolderManager.SUCCESS:
 			return status
 

managers/user.py

@@ -21,7 +21,7 @@
 import string, random, hashlib
 import uuid
 
-from db import User, session
+from db import User
 
 class UserManager:
 	SUCCESS = 0

scanner.py

@@ -27,11 +27,8 @@ def get_mime(ext):
 	return mimetypes.guess_type('dummy.' + ext, False)[0] or config.get('mimetypes', ext) or 'application/octet-stream'
 
 class Scanner:
-	def __init__(self, session):
-		self.__session = session
-		self.__tracks  = db.Track.query.all()
-		self.__artists = db.Artist.query.all()
-		self.__folders = db.Folder.query.all()
+	def __init__(self, store):
+		self.__store = store
 
 		self.__added_artists = 0
 		self.__added_albums  = 0