iTunes/iTunesParser.py

#!/usr/bin/env python

"""
--> Parse library and just do JSON

iTunes Graph Parser

Parses an iTunes library XML file and generates a JSON file
for use in the D3.js JavaScript library.

    Example Track info:
    {
        'Album': 'Nirvana',
        'Persistent ID': 'A50FE1436726815C',
        'Track Number': 4,
        'Location': 'file://localhost/Users/foo/Music/iTunes/iTunes%20Music/Nirvana/Nirvana/04%20Sliver.mp3',
        'File Folder Count': 4,
        'Album Rating Computed': True,
        'Total Time': 134295,
        'Sample Rate': 44100,
        'Genre': 'Rock/Alternative',
        'Bit Rate': 236,
        'Kind': 'MPEG audio file',
        'Name': 'Sliver',
        'Artist': 'Nirvana',
        'Date Added': datetime.datetime(2006, 10, 11, 4, 31, 38),
        'Album Rating': 60,
        'Rating': 40,
        'Date Modified': datetime.datetime(2009, 7, 18, 4, 57, 41),
        'Library Folder Count': 1,
        'Year': 2002,
        'Track ID': 7459,
        'Size': 3972838,
        'Track Type': 'File',
        'Play Count': 2,
        'Play Date UTC': datetime.datetime(2009, 7, 18, 5, 00, 00)
    }

"""

from __future__ import division
from optparse import OptionParser

import os
import io
import plistlib
import json
import datetime


class SetEncoder(json.JSONEncoder):
    def default(self, obj):
        if isinstance(obj, set):
            return list(obj)
        if isinstance(obj, datetime.datetime):
            return obj.isoformat()
            # encoded_object = int(mktime(obj.timetuple()))
        return json.JSONEncoder.default(self, obj)


class ITunesParser:
    def __init__(self, libraryFile):
        self._albums = {}
        self._artists = {}
        self.libraryFile = libraryFile

    def toJson(self):
        output = self._processSongs()

#        self._writeArtists()
#        self._writeAlbums()

        # return json.dumps(jsonObj, indent=indent, cls=SetEncoder)
        return output

    def toJsonP(self, rating=4):
        json = self.toJson()
        jsonp = ';itgCallback(' + json + ');'
        return jsonp


    def _readTracks(self):
        pl = plistlib.readPlist(self.libraryFile)
        return pl['Tracks']

    def _processSongs(self):
        tracks = self._readTracks()
        output = ""

        for k in tracks:
            track = tracks[k]

            # Filter out any non-music
            if track['Track Type'] != 'File':
                continue
            if 'Podcast' in track or 'Has Video' in track:
                continue

            persistentId = track['Persistent ID']
            jsonTrackIndic = {
                "index": {"_index": "itunessongs", "_type": "song", "_id": persistentId}
            }

            # Retrieve for each track artist information
            self._processArtist(track)
            # Retrieve for each track album information
            self._processAlbum(track)

            output += json.dumps(jsonTrackIndic, indent=None, cls=SetEncoder)
            output += "\n"
            output += json.dumps(track, indent=None, cls=SetEncoder)
            output += "\n"
        return output

    def _processArtist(self, track):
        if 'Artist' not in track:
            return

        akey = track['Artist']
        # Add artist
        if akey not in self._artists:
            self._artists[akey] = {
                'id': len(self._artists),
                'name': akey,
                'count': 0, 'plays': 0, 'rating': 0,
                'genres': set()
                }

        # Compute information
        rating = (track['Rating'] // 20) if 'Rating' in track else 0
        plays = track['Play Count'] if 'Play Count' in track else 0

        self._artists[akey]['count'] += 1
        self._artists[akey]['rating'] += rating
        self._artists[akey]['plays'] += plays

        if 'Genre' not in track:
            return

        # Split up the Genres
        genreParts = track['Genre'].split('/')
        self._artists[akey]['genres'] |= set(genreParts)
        return

    def _processAlbum(self, track):
        if 'Album' not in track:
            return

        akey = track['Album']
        if akey not in self._albums:
            self._albums[akey] = {
                'id': len(self._albums),
                'name': akey,
                'count': 0, 'plays': 0, 'rating': 0,
                'genres': set(),
                'artist': set()
            }

        # Compute information
        rating = (track['Rating'] // 20) if 'Rating' in track else 0
        plays = track['Play Count'] if 'Play Count' in track else 0

        self._albums[akey]['count'] += 1
        self._albums[akey]['rating'] += rating
        self._albums[akey]['plays'] += plays

        if 'Genre' not in track:
            return

        # Split up the Genres
        genreParts = track['Genre'].split('/')
        self._albums[akey]['genres'] |= set(genreParts)

        ## Add different artists
        if 'Artist' not in track:
            return
        self._albums[akey]['artist'].add(track['Artist'])
        return

#    def _writeArtists(self):
#        fileArtist = io.open('es-artist-data.json', 'wb')
#        for a in self._artists:
#            jsonTrackIndic = {
#                "index": {"_index": "itunessongs", "_type": "artist"}
#            }
#            fileArtist.write(json.dumps(jsonTrackIndic, indent=None, cls=SetEncoder))
#            fileArtist.write("\n")
#            fileArtist.write(json.dumps(self._artists[a], indent=None, cls=SetEncoder))
#            fileArtist.write("\n")
#        fileArtist.close()

#    def _writeAlbums(self):
#        fileAlbums = io.open('es-albums-data.json', 'wb')
#        for a in self._albums:
#            jsonTrackIndic = {
#                "index": {"_index": "itunessongs", "_type": "album"}
#            }
#            fileAlbums.write(json.dumps(jsonTrackIndic, indent=None, cls=SetEncoder))
#            fileAlbums.write("\n")
#            fileAlbums.write(json.dumps(self._albums[a], indent=None, cls=SetEncoder))
#            fileAlbums.write("\n")
#        fileAlbums.close()

#### main block ####

# Default input & output files
defaultLibraryFile = os.path.expanduser('iTunesMiniLibrary.xml')
defaultOutputFile = os.path.dirname(os.path.realpath(__file__)) + '/es-music-data.json'

# Get options
parser = OptionParser(version="%prog 1.0")
parser.add_option('-f', '--file', dest='file', type='string',
        help='iTunes Library XML file path',
        default=defaultLibraryFile)
parser.add_option('-o', '--output', dest='output', type='string',
        help='Output to file (default=./js/music-data.json)',
        default=defaultOutputFile)
parser.add_option('-c', '--console', dest='console', action='store_true',
        help='Output to console instead of file')
parser.add_option('-p', '--jsonp', dest='jsonp', action='store_true',
        help='Output in JSON-P format')
parser.add_option('-v', '--verbose', dest='verbose', action='store_true',
        help='Verbose output')

if __name__ == '__main__':
    (options, args) = parser.parse_args()

    itunesParser = ITunesParser(options.file)
    if options.jsonp:
        output = itunesParser.toJsonP()
    else:
        output = itunesParser.toJson()

    if options.console:
        print(output)
    else:
        with io.open(options.output, 'wb') as outfile:
            outfile.write(output)
        print('JSON data written to: ' + options.output)