diff options
Diffstat (limited to 'imdb-lookup')
-rwxr-xr-x | imdb-lookup/imdbinfo.py | 127 |
1 files changed, 127 insertions, 0 deletions
diff --git a/imdb-lookup/imdbinfo.py b/imdb-lookup/imdbinfo.py new file mode 100755 index 0000000..c4f80b5 --- /dev/null +++ b/imdb-lookup/imdbinfo.py @@ -0,0 +1,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)) |