From 06ea4e0754e11c574dd46ec6c0b79303f022609d Mon Sep 17 00:00:00 2001 From: Sarah Date: Sun, 3 Apr 2022 20:27:09 +0100 Subject: [PATCH] add keywords to data --- vcinema_utils/VCinemaUtils.py | 18 +++++++++++++++++- 1 file changed, 17 insertions(+), 1 deletion(-) diff --git a/vcinema_utils/VCinemaUtils.py b/vcinema_utils/VCinemaUtils.py index d5b2c94..06d7176 100644 --- a/vcinema_utils/VCinemaUtils.py +++ b/vcinema_utils/VCinemaUtils.py @@ -68,11 +68,27 @@ def add_imdb_data(imdb_id, viewings, data_fields, progressbar=None): progressbar.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): with ThreadPoolExecutor(4) as executor: future_imdb_tasks = set() - future_imdb_tasks.update(executor.submit(add_imdb_data, viewing['imdb_id'], viewings, field_names, progress_bar) 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 'keywords' in field_names: + future_imdb_tasks.update(executor.submit(add_imdb_keywords, viewing['imdb_id'], viewings, progress_bar) for viewing in viewings) progress_bar.max = len(future_imdb_tasks)