simpliefied caching

1 files changed, 101 insertions, 92 deletions

diff --git a/imdb-lookup/imdbinfo.py b/imdb-lookup/imdbinfo.py
index e1fe5be..7419103 100755
--- a/imdb-lookup/imdbinfo.py
+++ b/imdb-lookup/imdbinfo.py
@@ -5,7 +5,6 @@ import sys
 assert sys.version.startswith("3."), "Requires python3"
 import os
 import re
-import dbm
 import json
 import time
 import argparse
@@ -42,36 +41,62 @@ def read_key():
 
 
 class TMDBCache(object):
-    def __enter__(self):
+    def __init__(self):
         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_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):
+    def _get_cachedir(self, movie_id):
         if "XDG_CACHE_HOME" in os.environ.keys():
             cachedir = os.environ["XDG_CACHE_HOME"]
         else:
             cachedir = os.path.join(os.path.expanduser("~"), ".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)
+        cachedir = os.path.join(cachedir, "movies")
+        cachedir = os.path.join(cachedir, movie_id)
+        if not os.path.exists(cachedir):
+            os.makedirs(cachedir)
+
+        return cachedir
+
+    def _cache(self, movie_id, key, callable_func):
+        cachedir = self._get_cachedir(movie_id)
+        cachefile = os.path.join(cachedir, key)
+        if os.path.exists(cachefile):
+            fh = open(cachefile, "rb")
+            data = fh.read()
+            fh.close()
+            return data
+        else:
+            fh = open(cachefile, 'wb')
+            data = callable_func()
+            fh.write(data)
+            fh.close()
+            return data
+
+    def _cache_json(self, movie_id, key, callable_func):
+        def json_wrapper():
+            return json.dumps(callable_func()).encode('utf-8')
+        return json.loads(self._cache(movie_id, key, json_wrapper).decode('utf-8'))
+
+    def prune(self, movie_id):
+        self.logger.debug("prune {}".format(movie_id))
+        keys = [
+            "tmdb_info.json",
+            "tmdb_alternative_title.json",
+
+        ]
+        for key in keys:
+            cachefile = os.path.join(self._get_cachedir(movie_id), key)
+            if os.path.exists(cachefile):
+                self.logger.warn("Remove {}".format(cachefile))
+                os.unlink(cachefile)
+            else:
+                self.logger.debug("Key in db {}".format(cachefile))
 
     def infos(self, movie_id):
         try:
             self.logger.debug("movie %s", movie_id)
-            return self._cache(movie_id + "movies.info",
+            return self._cache_json(movie_id, "tmdb_info.json",
                                tmdb.Movies(movie_id).info)
         except Exception as e:
             raise Exception("Failed to query movie-id {id}: {reason}".format(
@@ -80,9 +105,9 @@ class TMDBCache(object):
     def alternative_title(self, movie_id, locale):
         """Returns the title in selected locale or original title otherwise"""
         try:
-            key = movie_id + "movies.alt_titles"
+            key = movie_id + "tmdb_alternative_title.json"
             search = tmdb.Movies(movie_id).alternative_titles
-            titles = self._cache(key, search)
+            titles = self._cache_json(movie_id, key, search)
             alt_title = list(filter(lambda l: l["iso_3166_1"] == locale,
                                     titles["titles"]))
             if alt_title:
@@ -94,42 +119,28 @@ class TMDBCache(object):
             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 = [
-            "imdb_maindetails_{}".format(movie_id),
-            movie_id + "movies.info",
-            movie_id + "movies.alt_titles",
-        ]
-        for key in keys:
-            if key in self.db:
-                self.logger.warn("Remove {}".format(key))
-                del self.db[key]
-            else:
-                self.logger.debug("Key not in db {}".format(key))
+    def poster(self, movie_id, f="w154"):
+        poster_path = self.infos(movie_id)['poster_path']
+        assert poster_path, "No poster_path defined for {}".format(movie_id)
 
-    def poster(self, poster_path, f="w154"):
-        self.logger.debug("poster %s", poster_path)
-        key = "poster_{}_{}".format(f, hashlib.md5(poster_path.encode('utf-8')).hexdigest()[0:10])
-        keyContentType = "{}_ct".format(key)
         url = "http://image.tmdb.org/t/p/{}/{}".format(f, poster_path)
-        if key not in self.db_images or keyContentType not in self.db_images:
-            r = requests.get(url)
-            self.db_images[key] = r.content
-            self.db_images[keyContentType] = r.headers['content-type']
+        self.logger.debug("poster %s", poster_path)
+        def download():
+            r = requests.get(url) # r.headers['content-type']
             r.close()
-        return (self.db_images[keyContentType], self.db_images[key])
+            return r.content
 
-    def poster_low(self, poster_path, format="w154"):
-        p = self.poster(poster_path, format)
-        if not p:
+        return self._cache(movie_id, "tmdb_poster.jpg", download)
+
+    def poster_low(self, movie_id, format="w154"):
+        data = self.poster(movie_id, format)
+        if not data:
             return None
-        contentType, data = p
         image = Image.open(BytesIO(data))
         image = image.filter(ImageFilter.GaussianBlur(radius=1))
         buf = BytesIO()
         image.save(buf, "JPEG", quality=18, optimize=True)
-        return ("image/jpeg", buf.getvalue())
+        return buf.getvalue()
 
     def _imdb_request(self, path, query):
         # see also https://github.com/richardasaurus/imdb-pie
@@ -163,8 +174,7 @@ class TMDBCache(object):
             r.close()
             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)
+        return self._cache_json(movie_id, "imdb_maindetails", do_request)
 
     def omdb_movie(self, movie_id):
         def do_request():
@@ -176,8 +186,7 @@ class TMDBCache(object):
             r.close()
             assert data is not None and data['Response'] == 'True', "Data must not be empty"
             return data
-        key = "omdb_{}".format(movie_id)
-        return self._cache(key, do_request)
+        return self._cache_json(movie_id, "omdb.json", do_request)
 
 
 def weight_rating(infos):
@@ -228,9 +237,9 @@ class Protector(object):
 
 
 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))
+    tmdbcache = TMDBCache()
+    for (filename, imdb_id) in imdb_ids:
+        print(tmdbcache.alternative_title(imdb_id, locale=args.lang))
 
 
 def do_data(args, imdb_ids):
@@ -248,34 +257,34 @@ def do_data(args, imdb_ids):
         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),
-                        "OMDB": tmdbcache.omdb_movie(imdb_id)}, sys.stdout)
+    tmdbcache = TMDBCache()
+    for (filename, imdb_id) in imdb_ids:
+        print_data({"TMDB": tmdbcache.infos(imdb_id),
+                    "IMDB": tmdbcache.imdb_movie(imdb_id),
+                    "OMDB": tmdbcache.omdb_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])
+    tmdbcache = 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)
+    tmdbcache = 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 = 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))
+    tmdbcache = TMDBCache()
+    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))
 
 
 def do_index(args, imdb_ids):
@@ -298,14 +307,15 @@ def do_index(args, imdb_ids):
         return out
 
     def data(callbackName):
-        db = Protector(tmdbcache)
-        def poster(imdb_id, data):
+        db = TMDBCache() #Protector(tmdbcache)
+        def poster(imdb_id):
+            data = db.poster_low(imdb_id)
             if data:
                 if not os.path.exists(".index.html/poster"):
                     os.makedirs(".index.html/poster")
 
                 out = os.path.join(".index.html/poster",imdb_id+".jpg")
-                open(out, "wb").write(data[1])
+                open(out, "wb").write(data)
                 return out
             else:
                 return None
@@ -323,7 +333,7 @@ def do_index(args, imdb_ids):
                 'id': imdb_id,
                 'title': tmdb['title'],
                 'path': path,
-                'poster': poster(imdb_id, db.poster_low(tmdb['poster_path'])),
+                'poster': poster(imdb_id),
                 'tagline': 'tagline' in imdb and imdb['tagline'] or None,
                 'plot': 'plot' in imdb and imdb['plot']['outline'] or None,
                 'website': 'homepage' in imdb and imdb['homepage'] or omdb['Website'] != 'N/A' and omdb['Website'] or None,
@@ -336,7 +346,6 @@ def do_index(args, imdb_ids):
                 'omdbUserTomato': (omdb['tomatoUserMeter'] != 'N/A') and float(omdb['tomatoUserMeter']) or None,
                 'omdbTomatoRating': (omdb['tomatoRating'] != 'N/A') and float(omdb['tomatoRating']) or None,
                 'omdbTomatoUserRating': (omdb['tomatoUserRating'] != 'N/A') and float(omdb['tomatoUserRating']) or None,
-                'omdbTomatoFresh': (omdb['tomatoFresh'] != 'N/A') and int(omdb['tomatoFresh']) or None,
                 'tmdbId': tmdb['id'],
             }
 
@@ -360,20 +369,20 @@ def do_index(args, imdb_ids):
     template_file = os.path.join(os.path.dirname(__file__),
                                  "index.jinja2.html")
     template = Template(open(template_file, "r").read())
-    with TMDBCache() as tmdbcache:
-        mapping = {
-            'title': 'Movie overview',
-            'install': install,
-            'data': data,
-        }
-        assert not os.path.exists("index.html") and not os.path.exists(".index.html"), \
-            "index.html or folder .index.html already exists"
-        stream = template.generate(mapping)
-        outfile = open("index.html", "wb")
-        for output in stream:
-            out = output.strip()
-            out = re.sub(" {2,}", " ", out)
-            outfile.write(out.encode('utf-8'))
+    tmdbcache = TMDBCache()
+    mapping = {
+        'title': 'Movie overview',
+        'install': install,
+        'data': data,
+    }
+    assert not os.path.exists("index.html") and not os.path.exists(".index.html"), \
+        "index.html or folder .index.html already exists"
+    stream = template.generate(mapping)
+    outfile = open("index.html", "wb")
+    for output in stream:
+        out = output.strip()
+        out = re.sub(" {2,}", " ", out)
+        outfile.write(out.encode('utf-8'))
 
 
 class HelpAction(argparse._HelpAction):