diff --git a/vcinema_utils/VCinemaUtils.py b/vcinema_utils/VCinemaUtils.py index a28a6c0..4570189 100644 --- a/vcinema_utils/VCinemaUtils.py +++ b/vcinema_utils/VCinemaUtils.py @@ -2,6 +2,7 @@ from collections import Counter from concurrent.futures import ThreadPoolExecutor, as_completed from progress.bar import Bar import sys +import functools from imdb_utils import IMDbUtils from bookstack import Bookstack @@ -49,18 +50,17 @@ def get_vcinema_viewings(token_id, token_secret, combine_repeat_viewings=True): return viewings -def get_imdb(imdb_id, bar): - imdb_entry = IMDbUtils.get_movie(imdb_id) - - if bar is not None: - bar.next() - - return imdb_entry +def increment_progressbar(bar, _): + bar.next() def add_imdb_data_to_viewings(viewings, field_names, progressbar=None): with ThreadPoolExecutor(4) as executor: - future_to_url = {executor.submit(get_imdb, viewing['imdb_id'], progressbar) for viewing in viewings} + future_to_url = {executor.submit(IMDbUtils.get_movie, viewing['imdb_id']) for viewing in viewings} + + if progressbar is not None: + for this_future in future_to_url: + this_future.add_done_callback(functools.partial(increment_progressbar, progressbar)) for future in as_completed(future_to_url): imdb_data = future.result()