summaryrefslogtreecommitdiff
path: root/imdb-lookup/imdbrating.py
blob: 662b8d9e8c17526deb2aefcbb1094c606e0344c4 (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
#!/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

if __name__ == u"__main__":
    tmdb.API_KEY = read_key()
    db = dbm.open(get_db_filename(), 'rw')

    for filename in sys.argv[1:]:
        imdb_id = filename.split('#')[-1]
        if imdb_id not in db.keys() or db[imdb_id] == 0:
            db[imdb_id] = str(get_rating(imdb_id))
        print u"{rating} {filename}".format(rating=db[imdb_id], filename=filename.decode('utf-8'))
    db.close()