path: root/imdb-lookup/imdbinfo.py

#!/usr/bin/env python2
# coding: utf-8
# python3 works as well
import os
import sys
import re
import dbm
import json
import base64
import argparse
import math
import logging
try:
    from urllib.parse import quote_plus
except:
    from urllib import quote_plus

try:
    import tmdbsimple as tmdb
    import requests
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.logger = logging.getLogger(self.__class__.__module__ + "." + self.__class__.__name__)
        self.logger.info("Open db")
        self.db = dbm.open(self._get_db_filename("tmdbmovie.dbm"),"c")
        self.db_images = dbm.open(self._get_db_filename("tmdbposter.dbm"), "c")
        return self

    def __exit__(self, type, value, traceback):
        self.logger.info("Close db")
        self.db.close()
        self.db_images.close()

    def _get_db_filename(self, name):
        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, name)

    def _cache(self, key, callable_func):
        if key not in self.db:
            self.db[key] = json.dumps(callable_func())
        d = self.db[key].decode('utf-8')
        return json.loads(d)

    def infos(self, movie_id):
        try:
            self.logger.debug("movie %s", movie_id)
            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 = list(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 poster(self, poster_path, format="w185"):
        self.logger.debug("poster %s", poster_path)
        key = "poster_{}_{}".format(format, poster_path)
        keyContentType = "poster_{}_{}_content_type".format(format, poster_path)
        url = "http://image.tmdb.org/t/p/{}/{}".format(format, poster_path)
        if key not in self.db_images:
            r = requests.get(url)
            self.db_images[key] = r.content
            self.db_images[keyContentType] = r.headers['content-type']
        return (self.db_images[keyContentType], self.db_images[key])

    def poster_base64(self, poster_path, format="w185"):
        p = self.poster(poster_path, format)
        if not p: return None
        contentType, data = p
        data64 = "".join(map(lambda c: isinstance(c,int) and chr(c) or c, 
                    filter(lambda c: c!='\n', base64.encodestring(data))))
        return "data:{};base64,{}".format(contentType, data64)
 

class Protector(object):
    def __init__(self, child):
        self.child = child
    def __getattr__(self, name):
        attr = getattr(self.child, name)
        def protected(*a, **kw):
            try:
                return attr(*a,**kw)
            except KeyboardInterrupt as e:
                raise e
            except:
                return None
        if callable(attr):
            return protected
        else:
            return o

def do_aka(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for (filename, 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 (filename, imdb_id) in imdb_ids:
            selected_properties = ["imdb_id", "revenue", "vote_average", "vote_count", "runtime", "budget", "vote_avarage", "release_date", "popularity", ]
            kv = map(lambda kv: u"{}={}".format(*kv), 
                    filter(lambda kv: kv[0] in selected_properties,
                            tmdbcache.infos(imdb_id).items()))
            print(u" ".join(kv))

def do_year(args, imdb_ids):
    with TMDBCache() as tmdbcache:
        for (filename, 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 (filename, imdb_id) in imdb_ids:
            tmdbcache.prune(imdb_id)

def do_rating(args, imdb_ids):
    """Calculates a rating based on vote_average and vote_count"""
    with TMDBCache() as tmdbcache:
        infos = list(filter(lambda i: "vote_average" in i and "vote_count" in i, 
                        map(lambda fid: tmdbcache.infos(fid[1]), imdb_ids)))
        maxvotes = max(map(lambda i: i["vote_count"], infos))
        for info in infos:
            f = math.sin(math.pi * ( info["vote_average"]/10.0 ) )
            d = (float(info["vote_count"]) / maxvotes) - 0.5
            info['rating'] = info["vote_average"] + 2 * d * f
            print(u"{rating:.02f} {imdb_id} {title:30s} avg={vote_average:.1f} count={vote_count:.0f}".format(**info))


def do_index(args, imdb_ids):
    """creates a index website"""
    def asBase64(poster):
        contentType, data = poster
        data64 = filter(lambda c: c!='\n', base64.encodestring(data))
        return "data:{};base64,{}".format(contentType, data64)
    try:
        from jinja2 import Template
    except ImportError:
        print("Failed to import jinja2 library for html-templating")
        sys.exit(1)
    template_file = os.path.join(os.path.dirname(__file__), "index.jinja2.html")
    template = Template(open(template_file).read())
    with TMDBCache() as tmdbcache:
        #infos = list(map(tmdbcache.infos, imdb_ids))
        #posters = list(map(asBase64, map(tmdbcache.poster, map(lambda info: info['poster_path'], infos))))
        mapping = {
            'input' : imdb_ids,
            'tmdbcache' : Protector(tmdbcache),
            'title' : 'Movie overview',
            'quote_plus' : quote_plus,
        }
        assert not os.path.exists("index.html"), "index.html already exists"
        stream = template.generate(mapping)
        outfile = open("index.html", "wb")
        for output in stream:
            output = output.encode("utf-8")
            outfile.write(output)

class HelpAction(argparse._HelpAction):
    def __call__(self, parser, namespace, values, option_string=None):
        formatter = parser._get_formatter()
        formatter.add_usage(parser.usage, parser._actions, parser._mutually_exclusive_groups)
        formatter.add_text(parser.description)

        formatter.start_section(parser._optionals.title)
        formatter.add_text(parser._optionals.description)
        formatter.add_arguments(parser._optionals._group_actions)
        formatter.end_section()
    
        subparsers_actions = [
            action for action in parser._actions
            if isinstance(action, argparse._SubParsersAction)]

        for subparsers_action in subparsers_actions:
            # get all subparsers and print help
            subparsers = subparsers_action.choices
            for subaction in subparsers_action._get_subactions():
                subparser = subparsers[subaction.dest]
                formatter.start_section("{} {} {}".format(formatter._prog, subaction.dest, 
                            formatter._format_actions_usage(subparser._actions, [])))
                formatter.add_text(subaction.help)
                formatter.add_arguments(subparser._positionals._group_actions)
                formatter.add_arguments(subparser._optionals._group_actions)
                formatter.end_section()

        print(formatter.format_help())
        parser.exit(0)

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

    parser = argparse.ArgumentParser(description="get movie data", add_help=False)
    parser.add_argument("--help", action=HelpAction, help="Display full help")
    parser.add_argument("--log-level", action='store', type=int, 
            help="Set log level (CRITICAL=50,ERROR=40,WARNING=30,INFO=20,DEBUG=10,NOTSET=0")
    parser.add_argument("-h", action=argparse._HelpAction, help="Display short help")

    subparsers = parser.add_subparsers()
    parser_aka = subparsers.add_parser("aka", add_help=False, help="Print alternative title in other languages")
    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="+", help="Files containing distinct movie-ids")

    parser_data = subparsers.add_parser("data", add_help=False, help="Print all available data")
    parser_data.set_defaults(func=do_data)
    parser_data.add_argument("files", action="append", nargs="+", help="Files containing distinct movie-ids")

    parser_year = subparsers.add_parser("year", add_help=False, help="Print only the release year")
    parser_year.set_defaults(func=do_year)
    parser_year.add_argument("files", action="append", nargs="+", help="Files containing distinct movie-ids")

    parser_prune = subparsers.add_parser("prune", add_help=False, help="Delete cache entries")
    parser_prune.set_defaults(func=do_prune)
    parser_prune.add_argument("files", action="append", nargs="+", help="Files containing distinct movie-ids")

    parser_rating = subparsers.add_parser("rating", add_help=False, help="Print movie ratings")
    parser_rating.set_defaults(func=do_rating)
    parser_rating.add_argument("files", action="append", nargs="+", help="Files containing distinct movie-ids")

    parser_index = subparsers.add_parser("index", add_help=False, help="Generate index.html file")
    parser_index.set_defaults(func=do_index)
    parser_index.add_argument("files", action="append", nargs="+", help="Files containing distinct movie-ids")

    args = parser.parse_args(sys.argv[1:])
    if "log_level" in args:
        logging.basicConfig(level=args.log_level)

    ids = map(lambda filename: (lambda x: (filename, 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))