diff options
Diffstat (limited to 'imdb-lookup')
| -rwxr-xr-x | imdb-lookup/imdbinfo.py | 184 |
1 files changed, 117 insertions, 67 deletions
diff --git a/imdb-lookup/imdbinfo.py b/imdb-lookup/imdbinfo.py index 8521e98..2f5e106 100755 --- a/imdb-lookup/imdbinfo.py +++ b/imdb-lookup/imdbinfo.py @@ -14,7 +14,6 @@ import math import logging import hashlib from io import BytesIO - from urllib.parse import quote as urlencode try: @@ -23,9 +22,11 @@ try: import requests except ImportError as e: print("Missing dependency: {0}".format(str(e))) - print("Install using system package manager or `pip install --user (one of: Pillow tmdbsimple requests)`") + print("Install using system package manager") + print("or `pip install --user (one of: Pillow tmdbsimple requests)`") sys.exit(1) + def read_key(): if "TMDB_KEY" in os.environ.keys(): return os.environ["TMDB_KEY"] @@ -36,14 +37,17 @@ def read_key(): if os.path.exists(os.path.join(cfg_home, "tmdbkey")): return open(os.path.join(cfg_home, "tmdbkey"), "r").read().strip() if os.path.exists(os.path.join(os.path.expanduser("~"), ".tmdbkey")): - return open(os.path.join(os.path.expanduser("~"), ".tmdbkey")).read().strip() - raise Exception("No TheMovieDB Key defined. Set Env. var. TMDB_KEY or .tmdbkey file") - + return open(os.path.join(os.path.expanduser("~"), + ".tmdbkey")).read().strip() + raise Exception("No TMDB Key defined. Set TMDB_KEY=.. or .tmdbkey file") + + class TMDBCache(object): def __enter__(self): - self.logger = logging.getLogger(self.__class__.__module__ + "." + self.__class__.__name__) + logger_name = self.__class__.__module__ + "." + self.__class__.__name__ + self.logger = logging.getLogger(logger_name) self.logger.info("Open db") - self.db = dbm.open(self._get_db_filename("tmdbmovie.dbm"),"c") + 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 @@ -68,23 +72,29 @@ class TMDBCache(object): 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) + 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))) + raise Exception("Failed to query movie-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: + key = movie_id + "movies.alt_titles" + search = tmdb.Movies(movie_id).search.alternative_titles + titles = self._cache(key, search) 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"])) + 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))) - + raise Exception("Failed to query movie-id {id}: {reason}".format( + id=movie_id, reason=str(e))) + def prune(self, movie_id): self.logger.debug("prune {}".format(movie_id)) keys = [ @@ -99,11 +109,11 @@ class TMDBCache(object): else: self.logger.debug("Key not in db {}".format(key)) - def poster(self, poster_path, format="w185"): + def poster(self, poster_path, f="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) + key = "poster_{}_{}".format(f, poster_path) + keyContentType = "poster_{}_{}_content_type".format(f, poster_path) + url = "http://image.tmdb.org/t/p/{}/{}".format(f, poster_path) if key not in self.db_images: r = requests.get(url) self.db_images[key] = r.content @@ -112,23 +122,27 @@ class TMDBCache(object): def poster_base64(self, poster_path, format="w185"): p = self.poster(poster_path, format) - if not p: return None + if not p: + return None contentType, data = p image = Image.open(BytesIO(data)) buf = BytesIO() image.save(buf, "JPEG", quality=10, optimize=True) - data64 = "".join(map(lambda c: isinstance(c,int) and chr(c) or c, - filter(lambda c: c!='\n', base64.encodestring(buf.getvalue())))) + data64 = base64.encodestring(buf.getvalue()) + data64 = "".join(map(lambda c: isinstance(c, int) and chr(c) or c, + filter(lambda c: c != '\n', data64))) return "data:{};base64,{}".format(contentType, data64) def _imdb_request(self, path, query): # see also https://github.com/richardasaurus/imdb-pie - # nice library but a bit strange API, so we chose to reimplement stuff here + # nice library but a bit strange API, + # so we chose to reimplement stuff here BASE_URI = 'app.imdb.com' API_KEY = '2wex6aeu6a8q9e49k7sfvufd6rhh0n' SHA1_KEY = hashlib.sha1(API_KEY.encode('utf8')).hexdigest() HEADERS = { - 'user-agent' : 'AppleWebKit/534.46 (KHTML, like Gecko) Ver sion/5.1 Mobile/9A405', + 'user-agent': 'AppleWebKit/534.46 (KHTML, like Gecko) ' + + 'Version/5.1 Mobile/9A405', } PARAMS = { "api": "v1", @@ -140,14 +154,15 @@ class TMDBCache(object): } q = query.copy() q.update(PARAMS) - return requests.get("https://{}{}".format(BASE_URI, path), params=q, headers=HEADERS) + return requests.get("https://{}{}".format(BASE_URI, path), + params=q, headers=HEADERS) def imdb_movie(self, movie_id): def do_request(): r = self._imdb_request("/title/maindetails", {'tconst': movie_id}) assert r.status_code == 200, "Request must return status-code 200" data = json.loads(r.text) - assert data != None and data['data'], "Data must not be empty" + assert data is not None and data['data'], "Data must not be empty" return data key = "imdb_maindetails_{}".format(movie_id) return self._cache(key, do_request) @@ -155,13 +170,15 @@ class TMDBCache(object): def weight_rating(infos): """ add 'rating' to all infos""" + infos = list(infos) maxvotes = max(map(lambda i: i["vote_count"], infos)) for info in infos: - f = math.sin(math.pi * ( info["vote_average"]/10.0 ) ) + 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 return infos + class Protector(object): """ The Protector saves the caller from exception. All callable attributes of child are dynamically @@ -180,11 +197,13 @@ 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) + return attr(*a, **kw) except KeyboardInterrupt as e: raise e except Exception as e: @@ -195,11 +214,13 @@ class Protector(object): else: return attr + 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): def print_data(data, io, indent=0): if isinstance(data, dict): @@ -209,67 +230,73 @@ def do_data(args, imdb_ids): else: io.write("\n{}{}: ".format(indent*" ", key)) print_data(val, io, indent+1) - elif isinstance(data, list) and len(data)>0: + elif isinstance(data, list) and len(data) > 0: for val in data: print_data(val, io, indent+2) - elif type(data) in (bool,str,int,float): + elif type(data) in (bool, str, int, float): io.write(str(data)) with TMDBCache() as tmdbcache: for (filename, imdb_id) in imdb_ids: - print_data({ - "TMDB" : tmdbcache.infos(imdb_id), - "IMDB" : tmdbcache.imdb_movie(imdb_id)}, sys.stdout) + print_data({"TMDB": tmdbcache.infos(imdb_id), + "IMDB": tmdbcache.imdb_movie(imdb_id)}, sys.stdout) + 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))) + infos = filter(lambda i: "vote_average" in i and "vote_count" in i, + map(lambda fid: tmdbcache.infos(fid[1]), imdb_ids)) weight_rating(infos) for info in infos: - print("{rating:.02f} {imdb_id} {title:30s} avg={vote_average:.1f} count={vote_count:.0f}".format(**info)) + print("{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""" + valid_extensions = [".mkv", ".avi", ".mov", ".mp4", ".mpg"] + def listMovieFiles(path): for root, dirs, files in os.walk(path): for curfile in files: - if sum(map(curfile.endswith, [".mkv", ".avi", ".mov", ".mp4"])): + if sum(map(curfile.endswith, valid_extensions)): yield os.path.join(root, curfile) def asBase64(poster): contentType, data = poster - data64 = filter(lambda c: c!='\n', base64.encodestring(data)) + 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_file = os.path.join(os.path.dirname(__file__), + "index.jinja2.html") template = Template(open(template_file, "r").read()) with TMDBCache() as tmdbcache: mapping = { - 'gmtime' : time.gmtime(), - 'input' : imdb_ids, - 'tmdbcache' : Protector(tmdbcache), - 'title' : 'Movie overview', - 'urlencode' : urlencode, - 'int' : int, - 'listMovieFiles' : listMovieFiles, - 'math' : math, + 'gmtime': time.gmtime(), + 'input': imdb_ids, + 'tmdbcache': Protector(tmdbcache), + 'title': 'Movie overview', + 'urlencode': urlencode, + 'int': int, + 'listMovieFiles': listMovieFiles, + 'math': math, } assert not os.path.exists("index.html"), "index.html already exists" stream = template.generate(mapping) @@ -277,16 +304,19 @@ def do_index(args, imdb_ids): for output in stream: outfile.write(output.encode('utf-8')) + 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_usage(parser.usage, + parser._actions, + parser._mutually_exclusive_groups) 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)] @@ -296,8 +326,10 @@ class HelpAction(argparse._HelpAction): 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, []))) + usage = formatter._format_actions_usage(subparser._actions, []) + formatter.start_section("{} {} {}".format(formatter._prog, + subaction.dest, + usage)) formatter.add_text(subaction.help) formatter.add_arguments(subparser._positionals._group_actions) formatter.add_arguments(subparser._optionals._group_actions) @@ -306,6 +338,7 @@ class HelpAction(argparse._HelpAction): print(formatter.format_help()) parser.exit(0) + def do_test(args): import doctest doctest.testmod() @@ -315,37 +348,52 @@ if __name__ == "__main__": parser = argparse.ArgumentParser(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") + 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 = subparsers.add_parser("aka", add_help=False, + help="Print alternative titles") 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_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 = 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_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 = 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_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 = 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_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 = 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_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 = 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") + parser_index.add_argument("files", action="append", nargs="+", + help="Files containing distinct movie-ids") - parser_test = subparsers.add_parser("test", add_help=False, help="Run testsuite") + parser_test = subparsers.add_parser("test", add_help=False, + help="Run testsuite") parser_test.set_defaults(func=do_test) args = parser.parse_args(sys.argv[1:]) @@ -353,11 +401,13 @@ if __name__ == "__main__": logging.basicConfig(level=args.log_level) if "files" in args: - ids = map(lambda filename: (lambda x: (filename, x.groups()[0]) if x else None)(re.match(".*#(tt[0-9]{7}).*", filename)), - args.files[0]) + ids = args.files[0] + ids = map(lambda filename: + (lambda x: (filename, x.groups()[0]) if x else None) + (re.match(".*#(tt[0-9]{7}).*", filename)), + ids) args.func(args, list(filter(lambda i: i is not None, ids))) elif "func" in args: args.func(args) else: HelpAction("")(parser, None, None) - |