diff --git a/get_subtitle.py b/get_subtitle.py index 949c06354..8ff8b4756 100644 --- a/get_subtitle.py +++ b/get_subtitle.py @@ -4,6 +4,7 @@ import os import sqlite3 import ast import logging +import operator import subprocess import time from datetime import datetime, timedelta @@ -28,21 +29,18 @@ def download_subtitle(path, language, hi, providers, providers_auth, sceneName, hi = True else: hi = False - if media_type == 'series': - type_of_score = 360 - minimum_score = float(get_general_settings()[8]) / 100 * type_of_score - elif media_type == 'movie': - type_of_score = 120 - minimum_score = float(get_general_settings()[22]) / 100 * type_of_score + language_set = set() + if language == 'pb': + language_set.add(Language('por', 'BR')) + else: + language_set.add(Language(language)) + use_scenename = get_general_settings()[9] + minimum_score = get_general_settings()[8] + minimum_score_movie = get_general_settings()[22] use_postprocessing = get_general_settings()[10] postprocessing_cmd = get_general_settings()[11] - if language == 'pob': - lang_obj = Language('por', 'BR') - else: - lang_obj = Language(language) - try: if sceneName is None or use_scenename is False: used_sceneName = False @@ -50,73 +48,102 @@ def download_subtitle(path, language, hi, providers, providers_auth, sceneName, else: used_sceneName = True video = Video.fromname(sceneName) + except: + logging.error("Error trying to get video information.") + else: + if media_type == "movie": + max_score = 120.0 + elif media_type == "series": + max_score = 360.0 + + try: + with AsyncProviderPool(max_workers=None, providers=providers, provider_configs=providers_auth) as p: + subtitles = p.list_subtitles(video, language_set) except Exception as e: - logging.exception('Error trying to extract information from this filename: ' + path) - return None + logging.exception("Error trying to get subtitle list from provider") else: - try: - best_subtitles = download_best_subtitles([video], {lang_obj}, providers=providers, min_score=minimum_score, hearing_impaired=hi, provider_configs=providers_auth) - except Exception as e: - logging.exception('Error trying to get the best subtitles for this file: ' + path) - return None - else: + subtitles_list = [] + #for s in subtitles: + # {s: compute_score(s, video, hearing_impaired=hi)} + sorted_subtitles = sorted([(s, compute_score(s, video, hearing_impaired=hi)) for s in subtitles], key=operator.itemgetter(1), reverse=True) + for s, preliminary_score in sorted_subtitles: + if media_type == "movie": + if (preliminary_score / max_score * 100) < int(minimum_score_movie): + continue + matched = set(s.get_matches(video)) + if hi == s.hearing_impaired: + matched.add('hearing_impaired') + not_matched = set(score.movie_scores.keys()) - matched + required = set(['title']) + if any(elem in required for elem in not_matched): + continue + elif media_type == "series": + if (preliminary_score / max_score * 100) < int(minimum_score): + continue + matched = set(s.get_matches(video)) + if hi == s.hearing_impaired: + matched.add('hearing_impaired') + not_matched = set(score.episode_scores.keys()) - matched + required = set(['series', 'season', 'episode']) + if any(elem in required for elem in not_matched): + continue + subtitles_list.append(s) + if len(subtitles_list) > 0: + best_subtitle = subtitles_list[0] + download_subtitles([best_subtitle], providers=providers, provider_configs=providers_auth) try: - best_subtitle = best_subtitles[video][0] - except: - logging.debug('No subtitles found for ' + path) + calculated_score = round(float(compute_score(best_subtitle, video, hearing_impaired=hi)) / max_score * 100, 2) + if used_sceneName == True: + video = scan_video(path) + single = get_general_settings()[7] + if single is True: + result = save_subtitles(video, [best_subtitle], single=True, encoding='utf-8') + else: + result = save_subtitles(video, [best_subtitle], encoding='utf-8') + except Exception as e: + logging.exception('Error saving subtitles file to disk.') return None else: - single = get_general_settings()[7] - try: - score = round(float(compute_score(best_subtitle, video, hearing_impaired=hi)) / type_of_score * 100, 2) - if used_sceneName == True: - video = scan_video(path) - if single is True: - result = save_subtitles(video, [best_subtitle], single=True, encoding='utf-8') - else: - result = save_subtitles(video, [best_subtitle], encoding='utf-8') - except: - logging.error('Error saving subtitles file to disk.') - return None + downloaded_provider = result[0].provider_name + downloaded_language = language_from_alpha3(result[0].language.alpha3) + downloaded_language_code2 = alpha2_from_alpha3(result[0].language.alpha3) + downloaded_language_code3 = result[0].language.alpha3 + downloaded_path = get_subtitle_path(path, language=language_set) + if used_sceneName == True: + message = downloaded_language + " subtitles downloaded from " + downloaded_provider + " with a score of " + unicode(calculated_score) + "% using this scene name: " + sceneName else: - downloaded_provider = result[0].provider_name - downloaded_language = language_from_alpha3(result[0].language.alpha3) - downloaded_language_code2 = alpha2_from_alpha3(result[0].language.alpha3) - downloaded_language_code3 = result[0].language.alpha3 - downloaded_path = get_subtitle_path(path, language=lang_obj) - if used_sceneName == True: - message = downloaded_language + " subtitles downloaded from " + downloaded_provider + " with a score of " + unicode(score) + "% using this scene name: " + sceneName - else: - message = downloaded_language + " subtitles downloaded from " + downloaded_provider + " with a score of " + unicode(score) + "% using filename guessing." - - if use_postprocessing is True: - command = pp_replace(postprocessing_cmd, path, downloaded_path, downloaded_language, downloaded_language_code2, downloaded_language_code3) - try: - if os.name == 'nt': - codepage = subprocess.Popen("chcp", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) - # wait for the process to terminate - out_codepage, err_codepage = codepage.communicate() - encoding = out_codepage.split(':')[-1].strip() - - process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + message = downloaded_language + " subtitles downloaded from " + downloaded_provider + " with a score of " + unicode(calculated_score) + "% using filename guessing." + + if use_postprocessing is True: + command = pp_replace(postprocessing_cmd, path, downloaded_path, downloaded_language, downloaded_language_code2, downloaded_language_code3) + try: + if os.name == 'nt': + codepage = subprocess.Popen("chcp", shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) # wait for the process to terminate - out, err = process.communicate() + out_codepage, err_codepage = codepage.communicate() + encoding = out_codepage.split(':')[-1].strip() - if os.name == 'nt': - out = out.decode(encoding) + process = subprocess.Popen(command, shell=True, stdout=subprocess.PIPE, stderr=subprocess.PIPE) + # wait for the process to terminate + out, err = process.communicate() + + if os.name == 'nt': + out = out.decode(encoding) - except: - if out == "": - logging.error('Post-processing result for file ' + path + ' : Nothing returned from command execution') - else: - logging.error('Post-processing result for file ' + path + ' : ' + out) + except: + if out == "": + logging.error('Post-processing result for file ' + path + ' : Nothing returned from command execution') else: - if out == "": - logging.info('Post-processing result for file ' + path + ' : Nothing returned from command execution') - else: - logging.info('Post-processing result for file ' + path + ' : ' + out) + logging.error('Post-processing result for file ' + path + ' : ' + out) + else: + if out == "": + logging.info('Post-processing result for file ' + path + ' : Nothing returned from command execution') + else: + logging.info('Post-processing result for file ' + path + ' : ' + out) - return message + return message + else: + return None def manual_search(path, language, hi, providers, providers_auth, sceneName, media_type): if hi == "True": @@ -129,11 +156,18 @@ def manual_search(path, language, hi, providers, providers_auth, sceneName, medi language_set.add(Language('por', 'BR')) else: language_set.add(Language(alpha3_from_alpha2(lang))) + + use_scenename = get_general_settings()[9] + use_postprocessing = get_general_settings()[10] + postprocessing_cmd = get_general_settings()[11] + try: - if sceneName != "None": - video = Video.fromname(sceneName) - else: + if sceneName is None or use_scenename is False: + used_sceneName = False video = scan_video(path) + else: + used_sceneName = True + video = Video.fromname(sceneName) except: logging.error("Error trying to get video information.") else: @@ -156,15 +190,21 @@ def manual_search(path, language, hi, providers, providers_auth, sceneName, medi if hi == s.hearing_impaired: matched.add('hearing_impaired') not_matched = set(score.movie_scores.keys()) - matched - if "title" in not_matched: + required = set(['title']) + if any(elem in required for elem in not_matched): continue + if used_sceneName: + not_matched.remove('hash') elif media_type == "series": matched = set(s.get_matches(video)) if hi == s.hearing_impaired: matched.add('hearing_impaired') not_matched = set(score.episode_scores.keys()) - matched - if "series" in not_matched or "season" in not_matched or "episode" in not_matched: + required = set(['series', 'season', 'episode']) + if any(elem in required for elem in not_matched): continue + if used_sceneName: + not_matched.remove('hash') subtitles_list.append(dict(score=round((compute_score(s, video, hearing_impaired=hi) / max_score * 100), 2), language=alpha2_from_alpha3(s.language.alpha3), hearing_impaired=str(s.hearing_impaired), provider=s.provider_name, subtitle=codecs.encode(pickle.dumps(s), "base64").decode(), url=s.page_link, matches=list(matched), dont_matches=list(not_matched))) subtitles_dict = {} subtitles_dict = sorted(subtitles_list, key=lambda x: x['score'], reverse=True)