summaryrefslogtreecommitdiff
path: root/imdb-lookup/imdbinfo.py
blob: c4f80b5a70404cd559acd9862faa766fd7d90c99 (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
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
#!/usr/bin/env python2
# coding: utf-8
# python3 works as well
import os
import sys
import re
import dbm
import json
import argparse

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")
        
class TMDBCache(object):
    def __enter__(self):
        self.db = dbm.open(self._get_db_filename(),"rw")
        return self

    def __exit__(self, type, value, traceback):
        self.db.close()

    def _get_db_filename(self):
        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, "tmdbmovie.dbm")

    def _cache(self, key, callable_func):
        if not self.db.has_key(key):
            self.db[key] = json.dumps(callable_func())
        return json.loads(self.db[key])

    def infos(self, movie_id):
        try:
            return self._cache(movie_id + "movies.info", tmdb.Movies(movie_id).info)
        except Exception as e:
            raise Exception("Failed to query movie with id {id}: {reason}".format(id=movie_id, reason=str(e)))

    def alternative_title(self, movie_id, locale):
        """Returns the title in selected locale or original title otherwise"""
        try:
            alt_title = filter(lambda l: l["iso_3166_1"] == locale,
                    self._cache(movie_id + "movies.alt_titles", tmdb.Movies(movie_id).alternative_titles)["titles"])
            if alt_title:
                return alt_title[0]["title"]
            else:
                infos = self.infos(movie_id)
                return infos["title"] or infos["original_title"]
        except Exception as e:
            raise Exception("Failed to query movie with id {id}: {reason}".format(id=movie_id, reason=str(e)))
    
    def prune(self, movie_id):
        keys = [movie_id + "movies.info", movie_id + "movies.alt_titles"]
        for key in keys:
            if key in self.db:
                print "Remove {}".format(key)
                del self.db[key]
        
def do_aka(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            print tmdbcache.alternative_title(imdb_id, locale=args.lang)

def do_data(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            selected_properties = ["imdb_id", "revenue", "vote_count", "runtime", "budget", "vote_avarage", "release_date", "popularity", ]
            kv = map(lambda kv: u"{}={}".format(*kv), 
                    filter(lambda (k,v): k in selected_properties,
                            tmdbcache.infos(imdb_id).items()))
            print u" ".join(kv)

def do_year(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            print tmdbcache.infos(imdb_id)["release_date"].split("-")[0]

def do_prune(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for imdb_id in imdb_ids:
            tmdbcache.prune(imdb_id)

if __name__ == u"__main__":
    tmdb.API_KEY = read_key()

    parser = argparse.ArgumentParser(description="get movie data")
    subparsers = parser.add_subparsers()
    parser_aka = subparsers.add_parser("aka")
    parser_aka.add_argument("--lang", help="Language code (default 'DE')")
    parser_aka.set_defaults(func=do_aka)
    parser_aka.add_argument("files", action="append", nargs="+")

    parser_data = subparsers.add_parser("data")
    parser_data.set_defaults(func=do_data)
    parser_data.add_argument("files", action="append", nargs="+")

    parser_year = subparsers.add_parser("year")
    parser_year.set_defaults(func=do_year)
    parser_year.add_argument("files", action="append", nargs="+")

    parser_prune = subparsers.add_parser("prune")
    parser_prune.set_defaults(func=do_prune)
    parser_prune.add_argument("files", action="append", nargs="+")

    args = parser.parse_args(sys.argv[1:])

    ids = map(lambda filename: (lambda x: x.groups()[0] if x else None)(re.match(".*#(tt[0-9]{7}).*", filename)),
                args.files[0])
    args.func(args, filter(lambda i: i is not None, ids))