Added support for multiple filenames

1 files changed, 10 insertions, 6 deletions

diff --git a/imdb-lookup/imdbrating.py b/imdb-lookup/imdbrating.py
index 995324c..dc2cd64 100755
--- a/imdb-lookup/imdbrating.py
+++ b/imdb-lookup/imdbrating.py
@@ -28,14 +28,18 @@ def read_key():
 
 
 def get_rating(imdb_id):
-    return tmdb.Find(id=imdb_id).info(external_source="imdb_id")['movie_results'][0]['vote_average']
+    info = tmdb.Find(id=imdb_id).info(external_source="imdb_id")
+    if 'movie_results' in info and len(info['movie_results']) == 1 and 'vote_average' in info['movie_results'][0]:
+        return info['movie_results'][0]['vote_average']
+    return 0
 
 if __name__ == u"__main__":
     tmdb.API_KEY = read_key()
-    imdb_id = sys.argv[1]
-
     db = dbm.open('cache', 'rw')
-    if imdb_id not in db.keys():
-        db[imdb_id] = str(get_rating(imdb_id))
-    print db[imdb_id]
+
+    for filename in sys.argv[1:]:
+        imdb_id = filename.split('#')[-1]
+        if imdb_id not in db.keys():
+            db[imdb_id] = str(get_rating(imdb_id))
+        print db[imdb_id]
     db.close()