summaryrefslogtreecommitdiff
path: root/imdb-lookup/imdbinfo.py
diff options
context:
space:
mode:
Diffstat (limited to 'imdb-lookup/imdbinfo.py')
-rwxr-xr-ximdb-lookup/imdbinfo.py127
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))