From 981bf3bb41780eafe5e9be958110d6ab3f029f8b Mon Sep 17 00:00:00 2001 From: Sarah Date: Sun, 3 Apr 2022 20:23:22 +0100 Subject: [PATCH] refactor --- vcinema_utils/VCinemaUtils.py | 46 ++++++++++++++++++----------------- 1 file changed, 24 insertions(+), 22 deletions(-) diff --git a/vcinema_utils/VCinemaUtils.py b/vcinema_utils/VCinemaUtils.py index 7d40a47..d5b2c94 100644 --- a/vcinema_utils/VCinemaUtils.py +++ b/vcinema_utils/VCinemaUtils.py @@ -1,7 +1,6 @@ from collections import Counter -from concurrent.futures import ThreadPoolExecutor, as_completed +from concurrent.futures import ThreadPoolExecutor import csv -import functools from imdb_utils import IMDbUtils from bookstack import Bookstack @@ -12,6 +11,7 @@ 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,32 +55,32 @@ def get_vcinema_viewings(token_id, token_secret, viewings_csv=None, combine_repe return viewings -def increment_progressbar(bar, _): - bar.next() +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 add_imdb_data_to_viewings(viewings, field_names, progressbar=None): +def add_imdb_data_to_viewings(viewings, field_names, progress_bar=None): with ThreadPoolExecutor(4) as executor: - future_imdb_tasks = {executor.submit(IMDbUtils.get_movie, viewing['imdb_id']) for viewing in viewings} + future_imdb_tasks = set() - if progressbar is not None: - for this_future in future_imdb_tasks: - this_future.add_done_callback(functools.partial(increment_progressbar, progressbar)) + future_imdb_tasks.update(executor.submit(add_imdb_data, viewing['imdb_id'], viewings, field_names, 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 progressbar is not None: - progressbar.finish() + if progress_bar is not None: + progress_bar.finish() -def filter_viewings(viewings, filter_field, progressbar=None): +def filter_viewings(viewings, filter_field, progress_bar=None): viewings_filtered = {} for viewing in viewings: @@ -98,7 +98,9 @@ def filter_viewings(viewings, filter_field, progressbar=None): else: viewings_filtered[viewing_field] = [viewing] - if progressbar is not None: - progressbar.next() + if progress_bar is not None: + progress_bar.next() + + progress_bar.finish() return viewings_filtered