imdbinfo

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))