get_artwork.py

#!/usr/bin/python
# -*- coding: utf-8 -*-

import os, sys
import urllib, urllib2
import json
import shutil

MOVIES = {
    'title_prompt': "Search for a movie title (or type 'tv' to switch to TV Shows): ",
    'search_url': "http://ax.itunes.apple.com/WebObjects/MZStoreServices.woa/wa/wsSearch?entity=movie&term=",
    'name_node': 'trackName',
}
TV = {
    'title_prompt': "Search for a TV Show season (or type 'movie' to switch to Movies): ",
    'search_url':  "http://ax.itunes.apple.com/WebObjects/MZStoreServices.woa/wa/wsSearch?entity=tvSeason&term=",
    'name_node': 'collectionName',
}

media = MOVIES
VIDEO_FILES = {'.mp4', '.m4v'}

def get_art(title, save_as):
    global media

    print "\nSearching for \"%s\"..." % title

    search_term = urllib.quote_plus(title)

    try:
        response = urllib2.urlopen("%s%s" % (media['search_url'], search_term))
        results = json.load(response)
        resultCount = results['resultCount']
        if resultCount > 0:
            which = None
            if resultCount > 1:
                for i, r in enumerate(results['results']):
                    if r['trackName'] == title or r['trackName'].startswith(title):
                        which = i
                        break
                else:
                    while True:
                        for index, result in enumerate(results['results']):
                            print "%s. %s" % (index+1, result[media['name_node']])
                        which = raw_input("\nEnter a number to download its artwork (or hit Enter to skip): ")
                        if not which:
                            return
                        which = int(which) - 1
                        if which < resultCount:
                            break
            else:
                which = 0
            orig_url = results['results'][which]['artworkUrl100']
            sys.stdout.write("Downloading artwork...")
            for url in [orig_url.replace("100x100-75.", ""), orig_url.replace("100x100-75.", "600x600-75.")]:
                try:
                    img_response = urllib2.urlopen(url)
                    with open(save_as, 'wb') as f:
                        shutil.copyfileobj(img_response, f)
                    sys.stdout.write(" done.\n\n")
                    break
                except urllib2.HTTPError:
                    pass
            else:
                print "Error fetching image for \"%s\"" % title
        else:
            print "No results found for \"%s\"" % title
    except urllib2.HTTPError:
        print "Error fetching image for \"%s\"" % title
        pass


def get_videos(path):
    for root, dirs, files in os.walk(path):
        for f in files:
            if os.path.splitext(f)[1].lower() in VIDEO_FILES:
                yield os.path.join(root, f)

if __name__ == "__main__":
    if len(sys.argv) < 2:
        print >>sys.stderr, "Usage: get_artwork.py <directory of video files>"
        sys.exit(1)
    for video in get_videos(sys.argv[1]):
        base_file = os.path.splitext(video)[0]
        title = os.path.basename(base_file)
        cover_file = base_file + '.jpg'
        if not os.path.isfile(cover_file):
            get_art(title, cover_file)