

class Track(object):
    """A Track object represents a single audio file"""
    def __init__(self, artist, title, album=None, filename=''):
        self.artist = artist
        self.title = title
        self.album = album
        self.filename = filename

    def __str__(self):
        s = "Track: %s: %s" % (self.artist, self.title)
        if self.filename:
            s += " [%s]" % self.filename
        return s
    
class Album(object):
    """An album object represents an album (a list of audio files)"""
    def __init__(self, artist, title, year='', genre='', tracks=[]):
        self.artist = artist
        self.title = title
        self.genre = genre
        self.tracks = tracks[:]  # could also use list(tracks)

    def add_track(self, track):
        self.tracks.append(track)

    def __str__(self):
        s = "Album: %s: %s" % (self.artist, self.title)
        if (len(self.tracks) > 0):
            s += " [%d tracks]" % len(self.tracks)
        return s

def music_library(tracks, albums):
    """Run the interactive music library application"""
    while True:
        c = raw_input("Search (s, sa, st, or S): ")
        (command,junk,query) = c.partition(' ')
        query = query.lower()
        if command == 's':
            print "Searching for '%s' in album titles..." % query, 
            results = [a for a in albums if query in a.title.lower()]
            print "done"
        elif command == 'sa':
            results = {t.artist for t in tracks if query in t.artist.lower()}
        elif command == 'st':
            results = [t for t in tracks if query in t.title.lower()]
        elif command == 'S':
            results = [t for t in tracks \
                        if query in t.title.lower() \
                        or query in t.artist.lower()]
            results.extend([a for a in albums \
                        if query in a.title.lower() \
                        or query in a.artist.lower() \
                        or query in a.genre.lower()])
        else:
            print "Search command must be one of s, sa, st, or S"
            results = []
        for x in results: print "  %s" % str(x)

import os
from mutagen.mp3 import MP3

def load_library(dir):
    """Load a music library from a directory"""
    tracks = []
    albums = {}
    for root, dirs, files in os.walk(dir):
        for filename in files:
            if filename.lower().endswith(".mp3"):
                fullname = "%s/%s" % (root, filename)
                try:
                    audio = MP3(fullname)
                except EOFError:
                    continue
                artist,title,albumtitle,genre,year = 5*('',)
                if 'TPE1' in audio: artist = str(audio['TPE1'])
                if 'TIT2' in audio: title = str(audio['TIT2'])
                if 'TALB' in audio: albumtitle = str(audio['TALB'])
                if albumtitle in albums:
                    album = albums[albumtitle]
                else:
                    if 'TCON' in audio: genre = str(audio['TCON'])
                    if 'TDRC' in audio: year = str(audio['TDRC'])
                    album = Album(artist, albumtitle, year, genre)
                    albums[albumtitle] = album
                track = Track(artist, title, album)
                tracks.append(track)
                album.add_track(track)
    return (tracks,[albums[a] for a in albums])


