vcinema_utils/VCinemaUtils.py

@@ -1,6 +1,7 @@
 from collections import Counter
-from concurrent.futures import ThreadPoolExecutor
+from concurrent.futures import ThreadPoolExecutor, as_completed
 import csv
+import functools
 
 from imdb_utils import IMDbUtils
 from bookstack import Bookstack
@@ -11,7 +12,6 @@ JACKNET_WIKI_URL = "https://wiki.jacknet.io"
 # Page ID of https://wiki.jacknet.io/books/vcinema/page/csv
 CSV_PAGE_ID = 11
 
-
 def get_viewings_csv_attachment_id(token_id, token_secret):
     attachments = Bookstack.get_attachments(JACKNET_WIKI_URL, token_id, token_secret)
 
@@ -55,48 +55,32 @@ def get_vcinema_viewings(token_id, token_secret, viewings_csv=None, combine_repe
     return viewings
 
 
-def add_imdb_data(imdb_id, viewings, data_fields, progressbar=None):
-    movie = IMDbUtils.get_movie(imdb_id)
-
-    for viewing in viewings:
-        if viewing['imdb_id'] == movie.movieID:
-            for field_name in data_fields:
-                if field_name in movie:
-                    viewing[field_name] = movie[field_name]
-
-    if progressbar is not None:
-        progressbar.next()
+def increment_progressbar(bar, _):
+    bar.next()
 
 
-def add_imdb_keywords(imdb_id, viewings, progressbar=None):
-    movie = IMDbUtils.get_movie_keywords(imdb_id)
-
-    for viewing in viewings:
-        if viewing['imdb_id'] == movie.movieID:
-            if 'keywords' in movie:
-                viewing['keywords'] = movie['keywords']
-
-    if progressbar is not None:
-        progressbar.next()
-
-
-def add_imdb_data_to_viewings(viewings, field_names, progress_bar=None):
+def add_imdb_data_to_viewings(viewings, field_names, progressbar=None):
     with ThreadPoolExecutor(4) as executor:
-        future_imdb_tasks = set()
+        future_imdb_tasks = {executor.submit(IMDbUtils.get_movie, viewing['imdb_id']) for viewing in viewings}
 
-        if ('keywords' in field_names and field_names > 1) or ('keywords' not in field_names and field_names > 0):
-            future_imdb_tasks.update(executor.submit(add_imdb_data, viewing['imdb_id'], viewings, field_names, progress_bar) for viewing in viewings)
+        if progressbar is not None:
+            for this_future in future_imdb_tasks:
+                this_future.add_done_callback(functools.partial(increment_progressbar, progressbar))
 
-        if 'keywords' in field_names:
-            future_imdb_tasks.update(executor.submit(add_imdb_keywords, viewing['imdb_id'], viewings, progress_bar) for viewing in viewings)
+        for future in as_completed(future_imdb_tasks):
+            imdb_data = future.result()
 
-        progress_bar.max = len(future_imdb_tasks)
+            for viewing in viewings:
+                if viewing['imdb_id'] == imdb_data.movieID:
+                    for field_name in field_names:
+                        if field_name in imdb_data:
+                            viewing[field_name] = imdb_data[field_name]
 
-    if progress_bar is not None:
-        progress_bar.finish()
+        if progressbar is not None:
+            progressbar.finish()
 
 
-def filter_viewings(viewings, filter_field, progress_bar=None):
+def filter_viewings(viewings, filter_field, progressbar=None):
     viewings_filtered = {}
 
     for viewing in viewings:
@@ -114,9 +98,7 @@ def filter_viewings(viewings, filter_field, progress_bar=None):
                 else:
                     viewings_filtered[viewing_field] = [viewing]
 
-        if progress_bar is not None:
-            progress_bar.next()
-
-    progress_bar.finish()
+        if progressbar is not None:
+            progressbar.next()
 
     return viewings_filtered