This commit is contained in:
Sarah 2022-04-03 20:23:22 +01:00
parent c7f4b47da0
commit 981bf3bb41

View File

@ -1,7 +1,6 @@
from collections import Counter from collections import Counter
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor
import csv import csv
import functools
from imdb_utils import IMDbUtils from imdb_utils import IMDbUtils
from bookstack import Bookstack 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 # Page ID of https://wiki.jacknet.io/books/vcinema/page/csv
CSV_PAGE_ID = 11 CSV_PAGE_ID = 11
def get_viewings_csv_attachment_id(token_id, token_secret): def get_viewings_csv_attachment_id(token_id, token_secret):
attachments = Bookstack.get_attachments(JACKNET_WIKI_URL, 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 return viewings
def increment_progressbar(bar, _): def add_imdb_data(imdb_id, viewings, data_fields, progressbar=None):
bar.next() movie = IMDbUtils.get_movie(imdb_id)
def add_imdb_data_to_viewings(viewings, field_names, progressbar=None):
with ThreadPoolExecutor(4) as executor:
future_imdb_tasks = {executor.submit(IMDbUtils.get_movie, viewing['imdb_id']) 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))
for future in as_completed(future_imdb_tasks):
imdb_data = future.result()
for viewing in viewings: for viewing in viewings:
if viewing['imdb_id'] == imdb_data.movieID: if viewing['imdb_id'] == movie.movieID:
for field_name in field_names: for field_name in data_fields:
if field_name in imdb_data: if field_name in movie:
viewing[field_name] = imdb_data[field_name] viewing[field_name] = movie[field_name]
if progressbar is not None: if progressbar is not None:
progressbar.finish() progressbar.next()
def filter_viewings(viewings, filter_field, progressbar=None): 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)
progress_bar.max = len(future_imdb_tasks)
if progress_bar is not None:
progress_bar.finish()
def filter_viewings(viewings, filter_field, progress_bar=None):
viewings_filtered = {} viewings_filtered = {}
for viewing in viewings: for viewing in viewings:
@ -98,7 +98,9 @@ def filter_viewings(viewings, filter_field, progressbar=None):
else: else:
viewings_filtered[viewing_field] = [viewing] viewings_filtered[viewing_field] = [viewing]
if progressbar is not None: if progress_bar is not None:
progressbar.next() progress_bar.next()
progress_bar.finish()
return viewings_filtered return viewings_filtered