use guessit

1 files changed, 46 insertions, 72 deletions

diff --git a/imdb-lookup/imdbinfo.py b/imdb-lookup/imdbinfo.py
index b2691cd..ca0d7b1 100755
--- a/imdb-lookup/imdbinfo.py
+++ b/imdb-lookup/imdbinfo.py
@@ -14,18 +14,14 @@ import shutil
 import logging
 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
-try:
-    from PIL import Image, ImageFilter
-    import tmdbsimple as tmdb
-    import requests
-except ImportError as e:
-    print("Missing dependency: {0}".format(str(e)))
-    print("Install using system package manager")
-    print("or `pip install --user (one of: Pillow tmdbsimple requests)`")
-    sys.exit(1)
+from PIL import Image, ImageFilter
+import tmdbsimple as tmdb
+import requests
+import guessit
 
 def read_key():
     if "TMDB_KEY" in os.environ.keys():
@@ -206,34 +202,13 @@ def weight_rating(infos):
     return infos
 
 def filter_filename(filename):
-    """ remove possibly everything except the movie name
-    >>> filter_filename("The.Interview.2014.1080p")
-    'The Interview'
-    >>> filter_filename("Hello World.webm")
-    'Hello World'
-    >>> filter_filename("Foo bar Blah ldbla (2014) #tt2123123")
-    'Foo bar Blah ldbla'
-    >>> filter_filename("When.Bla.Foo.1948.German.1080p.BluRay.x264-Foobar")
-    'When Bla Foo'
-    >>> filter_filename("St.Franziskus.2008.WEBRip.MD.German.x264-NSane")
-    'St Franziskus'
-    """
-    stopwords = [ 
-        "dvd", "ac3", "r5", "unrated", "720p", "md",
-        "ts", "ld", "bdrip", "tvrip", "dvdrip", "dvdscr", "uncut",
-        "german", "english", "telesync", "x264", "hdtv", "ws", "webm",
-        "20[0-1][0-9]", "19[0-9][0-9]", "\\(20[0-1][0-9]\\)", "\\(19[0-9][0-9]\\)",
-    ]   
-    def findword(word):
-        return re.findall("(.*?)[\. ]" + word + "($|[\.\- ])", filename, re.IGNORECASE)
-    matches = [[findword(word), word] for word in stopwords]
-    matches = [match for match in matches if match[0] != []]
-    matches.sort() # make the shortest the first in list
-    if len(matches) > 0:
-        name = matches[0][0][0][0]
+    """remove possibly everything except the movie name"""
+    info = guessit.guess_movie_info(filename).to_dict()
+    print("Parsed: {}".format(", ".join(map(lambda kv: "{}={}".format(*kv), info.items()))))
+    if "title" in info:
+        return info["title"]
     else:
-        name = filename
-    return name.replace(u".", u" ").strip()
+        return os.path.basename(filename).replace(".", " ")
 
 class Protector(object):
     """ The Protector saves the caller from exception.
@@ -458,42 +433,41 @@ def do_rename(args):
         """remove potentially harmful characters from filename"""
         return re.sub(u"[^a-zA-Z0-9()\- _#]", u"_", filename)
 
-    from pprint import pprint
-    countrycode = locale.getdefaultlocale()[0][3:5]
-    db = Protector(TMDBCache())
-    search_term = filter_filename(args.file[0])
-    search_term = input("Search [{}]: ".format(search_term)) or search_term
-    results = db.tmdb_search(search_term)
-    results.sort(key=lambda i:i["popularity"],reverse=True)
-    results = filter(lambda i: i["imdb_id"] and i["release_date"], results)
-    results = list(results)
-    
-    for (i,result) in zip(range(len(results)), results):
-        result["i"] = i+1
-        result["year"] = result["release_date"].split("-")[0]
-        print("{i} {title} ({year}) #{imdb_id}".format(**result))
-        result["alternative_title"] = db.alternative_title(result["imdb_id"], locale=countrycode)
-        if result["alternative_title"] != result["title"]:
-            print("  a.k.a.: {}".format(result["alternative_title"]))
-
-    select = None
-    while not select:
-        try:
-            select = int(input("Choose: "))
-        except ValueError:
-            pass
-
-    filename = "{alternative_title} ({year}) #{imdb_id}".format(**results[select-1])
-    filename = substitute_filename(filename)
-    
-    if os.path.exists(args.file[0]):
-        choice = input("Move '{}' -> '{}' ? [y/N]: ".format(args.file[0], filename)).lower()
-        if choice == "y":
+    for input_file in args.file:
+        countrycode = locale.getdefaultlocale()[0][3:5]
+        db = Protector(TMDBCache())
+        search_term = filter_filename(input_file)
+        search_term = input("Search [{}]: ".format(search_term)) or search_term
+        results = db.tmdb_search(search_term)
+        results.sort(key=lambda i:i["popularity"],reverse=True)
+        results = filter(lambda i: i["imdb_id"] and i["release_date"], results)
+        results = list(results)
+        
+        select = None
+        while not select:
+            for (i,result) in zip(range(len(results)), results):
+                result["i"] = i+1
+                result["year"] = result["release_date"].split("-")[0]
+                print("{i} {title} ({year}) - check: http://www.imdb.com/title/{imdb_id}".format(**result))
+                result["alternative_title"] = db.alternative_title(result["imdb_id"], locale=countrycode)
+                if result["alternative_title"] != result["title"]:
+                    print("  a.k.a.: {}".format(result["alternative_title"]))
+            try:
+                select = int(input("Choose: "))
+            except ValueError:
+                pass
+
+        filename = "{alternative_title} ({year}) #{imdb_id}".format(**results[select-1])
+        filename = substitute_filename(filename)
+        
+        if os.path.exists(input_file):
+            choice = input("Move '{}' -> '{}' ? [y/N]: ".format(input_file, filename)).lower()
+            if choice == "y":
+                os.mkdir(filename)
+                shutil.move(input_file, filename)
+        else:
+            print("mkdir '{}'".format(filename))
             os.mkdir(filename)
-            shutil.move(args.file[0], filename)
-    else:
-        print("mkdir '{}'".format(filename))
-        os.mkdir(filename)
 
 def do_test(args):
     import doctest
@@ -551,7 +525,7 @@ if __name__ == "__main__":
     parser_rename = subparsers.add_parser("rename", add_help=False,
                               help="Search supported rename")
     parser_rename.set_defaults(func=do_rename)
-    parser_rename.add_argument("file", metavar="PATH", nargs=1,
+    parser_rename.add_argument("file", metavar="PATH", nargs="*",
                               help="File or folder to rename into tagged-imdb format")
 
     parser_test = subparsers.add_parser("test", add_help=False,