summaryrefslogtreecommitdiff
path: root/imdb-lookup/imdbrating.py
blob: da40ed53b95c057d1c3295da6178be339ee956ef (plain)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
#!/usr/bin/env python2
# coding: utf-8
# python3 works as well
import os
import sys
import shutil
import dbm

try:
    import tmdbsimple as tmdb
except ImportError as e:
    print(u"Missing dependency: {0}".format(str(e)))
    print(u"Install using system package manager or `pip install --user <module>`")
    sys.exit(1)

def read_key():
    if u"TMDB_KEY" in os.environ.keys():
        return os.environ[u"TMDB_KEY"]
    if u"XDG_CONFIG_HOME" in os.environ.keys():
        cfg_home = os.environ[u"XDG_CONFIG_HOME"]
    else:
        cfg_home = os.path.join(os.path.expanduser(u"~"), ".config")
    if os.path.exists(os.path.join(cfg_home, u"tmdbkey")):
        return open(os.path.join(cfg_home, u"tmdbkey"), "r").read().strip()
    if os.path.exists(os.path.join(os.path.expanduser(u"~"), ".tmdbkey")):
        return open(os.path.join(os.path.expanduser(u"~"), ".tmdbkey")).read().strip()
    raise Exception(u"No TheMovieDB Key defined. Set Env. var. TMDB_KEY or .tmdbkey file")
        
def get_db_filename():
        if u"XDG_CACHE_HOME" in os.environ.keys():
            cachedir = os.environ["XDG_CACHE_HOME"]
        else:
            cachedir = os.path.join(os.path.expanduser(u"~"), ".cache")
        return os.path.join(cachedir, "imdbrating.dbm")

def get_rating(imdb_id):
    info = tmdb.Find(id=imdb_id).info(external_source="imdb_id")
    if 'movie_results' in info and len(info['movie_results']) == 1 and 'vote_average' in info['movie_results'][0]:
        return info['movie_results'][0]['vote_average']
    return 0

def print_ratings(files, db):
    for filename in files:
        if os.path.basename(filename).startswith('.'):
            continue
        elif os.path.isdir(filename):
            print_ratings(map(lambda p: os.path.join(filename, p), os.listdir(filename)), db)
        imdb_id = os.path.basename(filename).split('#')[-1]
        if imdb_id.startswith('tt'):
            if imdb_id not in db.keys() or db[imdb_id] == 0:
                db[imdb_id] = str(get_rating(imdb_id))
            print "{rating} {filename}".format(rating=db[imdb_id], filename=filename)

if __name__ == u"__main__":
    tmdb.API_KEY = read_key()
    db = dbm.open(get_db_filename(), 'rw')
    files = sum([os.listdir(x) if os.path.isdir(x) else [x] for x in sys.argv[1:]],[])
    print_ratings(files, db)
    db.close()