guessit for tags in index

1 files changed, 16 insertions, 17 deletions

diff --git a/imdb-lookup/imdbinfo.py b/imdb-lookup/imdbinfo.py
index ca0d7b1..84845ce 100755
--- a/imdb-lookup/imdbinfo.py
+++ b/imdb-lookup/imdbinfo.py
@@ -16,7 +16,6 @@ import hashlib
 from io import BytesIO
 from pprint import pprint
 from multiprocessing.dummy import Pool
-from threading import Lock
 from urllib.parse import quote as urlencode
 from PIL import Image, ImageFilter
 import tmdbsimple as tmdb
@@ -175,6 +174,7 @@ class TMDBCache(object):
             data = json.loads(r.text)
             r.close()
             assert data is not None and data['data'], "Data must not be empty"
+            assert 'rating' in data['data']
             return data['data']
         return self._cache_json(movie_id, "imdb_maindetails.json", do_request)
 
@@ -238,7 +238,7 @@ class Protector(object):
             except KeyboardInterrupt as e:
                 raise e
             except Exception as e:
-                logging.error("Error calling %s: %s: %s", name, type(e), e)
+                logging.error("Error calling %s: '%s': '%s' with args '(%s,%s)'", name, type(e), e, a ,kw)
                 return None
         if callable(attr):
             return protected
@@ -321,19 +321,14 @@ def do_index(args, imdb_ids):
             def find():
                 for root, dirs, files in os.walk(path):
                     for curfile in files:
-                        if curfile.endswith("mkv"):
-                            yield "MKV"
-                        if curfile.endswith("avi"):
-                            yield "AVI"
-                    for curdir in dirs:
-                        tags = { "german" : "German", 
-                                 "1080" : "1080p", "720" : "720p",
-                                 "bdrip" : "BlueRay", "blueray" : "BlueRay",
-                                 "hevc" : "HEVC" }
-                        for (key, tag) in tags.items():
-                            if key in curdir.lower():
-                                yield tag
-
+                        if curfile.lower()[-3:] in ["mkv","avi","mp4"]:
+                            info = guessit.guess_movie_info(os.path.join(root,curfile)).to_dict()
+                            for key,val in info.items():
+                                if key in ["format","videoCodec","releaseGroup","container","container"]:
+                                    if not isinstance(val, list):
+                                        val = [val]
+                                    for lval in val:
+                                        yield "{}:{}".format(key,lval)
             l = list(set(find()))
             l.sort()
             return l
@@ -366,7 +361,8 @@ def do_index(args, imdb_ids):
             print("Error in {} {}".format(path, imdb_id))
             return None
         else:
-            return {
+            try: 
+                return {
             'id': imdb_id,
             'title': tmdb['title'],
             'summary': compile_summary(imdb, omdb, tmdb),
@@ -383,7 +379,10 @@ def do_index(args, imdb_ids):
             'omdbTomatoMeter': (omdb['tomatoMeter'] != 'N/A') and float(omdb['tomatoMeter']) or None,
             'omdbTomatoRating': (omdb['tomatoRating'] != 'N/A') and float(omdb['tomatoRating']) or None,
             'tmdbId': tmdb['id'],
-        }
+                }
+            except Exception as e:
+                print("Error in {} {}: {}".format(path, imdb_id, str(e)))
+                return None
 
     datadir = os.path.join(os.path.dirname(__file__), "html")
     index_files = os.path.join(datadir, "index-files")