forked from morpheus65535/bazarr
-
Notifications
You must be signed in to change notification settings - Fork 0
/
get_subtitle.py
480 lines (425 loc) · 23.2 KB
/
get_subtitle.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
from get_argv import config_dir
import os
import sqlite3
import ast
import logging
import operator
import subprocess
import time
from datetime import datetime, timedelta
from babelfish import Language
from subliminal import region, scan_video, Video, download_best_subtitles, compute_score, save_subtitles, AsyncProviderPool, score, list_subtitles, download_subtitles
from subliminal.subtitle import get_subtitle_path
from get_languages import language_from_alpha3, alpha2_from_alpha3, alpha3_from_alpha2
from bs4 import UnicodeDammit
from get_settings import get_general_settings, pp_replace, path_replace, path_replace_movie, path_replace_reverse, path_replace_reverse_movie
from list_subtitles import store_subtitles, list_missing_subtitles, store_subtitles_movie, list_missing_subtitles_movies
from utils import history_log, history_log_movie
from notifier import send_notifications, send_notifications_movie
import cPickle as pickle
import codecs
from get_providers import get_providers, get_providers_auth
# configure the cache
region.configure('dogpile.cache.memory')
def download_subtitle(path, language, hi, providers, providers_auth, sceneName, media_type):
if hi == "True":
hi = True
else:
hi = False
language_set = set()
if language == 'pob':
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]
try:
if sceneName == "None" or use_scenename is False:
used_sceneName = False
video = scan_video(path)
else:
used_sceneName = True
video = Video.fromname(sceneName)
except Exception as e:
logging.exception("Error trying to get video information for this file: " + path)
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 get subtitle list from provider")
else:
subtitles_list = []
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:
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:
if len(result) > 0:
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:
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_codepage, err_codepage = codepage.communicate()
encoding = out_codepage.split(':')[-1].strip()
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)
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
else:
return None
else:
return None
def manual_search(path, language, hi, providers, providers_auth, sceneName, media_type):
if hi == "True":
hi = True
else:
hi = False
language_set = set()
for lang in ast.literal_eval(language):
lang = alpha3_from_alpha2(lang)
if lang == 'pob':
language_set.add(Language('por', 'BR'))
else:
language_set.add(Language(lang))
use_scenename = get_general_settings()[9]
use_postprocessing = get_general_settings()[10]
postprocessing_cmd = get_general_settings()[11]
try:
if sceneName == "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:
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 get subtitle list from provider")
else:
subtitles_list = []
for s in subtitles:
{s: compute_score(s, video, hearing_impaired=hi)}
if media_type == "movie":
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
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
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)
return(subtitles_dict)
def manual_download_subtitle(path, language, hi, subtitle, provider, providers_auth, sceneName, media_type):
if hi == "True":
hi = True
else:
hi = False
subtitle = pickle.loads(codecs.decode(subtitle.encode(), "base64"))
if media_type == 'series':
type_of_score = 360
elif media_type == 'movie':
type_of_score = 120
use_scenename = get_general_settings()[9]
use_postprocessing = get_general_settings()[10]
postprocessing_cmd = get_general_settings()[11]
language = alpha3_from_alpha2(language)
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
video = scan_video(path)
else:
used_sceneName = True
video = Video.fromname(sceneName)
except Exception as e:
logging.exception('Error trying to extract information from this filename: ' + path)
return None
else:
try:
best_subtitle = subtitle
download_subtitles([best_subtitle], providers=provider, provider_configs=providers_auth)
except Exception as e:
logging.exception('Error downloading subtitles for ' + path)
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 Exception as e:
logging.exception('Error saving subtitles file to disk.')
return None
else:
if len(result) > 0:
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)
message = downloaded_language + " subtitles downloaded from " + downloaded_provider + " with a score of " + unicode(score) + "% using manual search."
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)
# 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)
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
else:
return None
def series_download_subtitles(no):
if get_general_settings()[24] is True:
monitored_only_query_string = ' AND monitored = "True"'
else:
monitored_only_query_string = ""
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
episodes_details = c_db.execute('SELECT path, missing_subtitles, sonarrEpisodeId, scene_name FROM table_episodes WHERE sonarrSeriesId = ? AND missing_subtitles != "[]"' + monitored_only_query_string, (no,)).fetchall()
series_details = c_db.execute("SELECT hearing_impaired FROM table_shows WHERE sonarrSeriesId = ?", (no,)).fetchone()
c_db.close()
providers_list = get_providers()
providers_auth = get_providers_auth()
for episode in episodes_details:
for language in ast.literal_eval(episode[1]):
if language is not None:
message = download_subtitle(path_replace(episode[0]), str(alpha3_from_alpha2(language)), series_details[0], providers_list, providers_auth, str(episode[3]), 'series')
if message is not None:
store_subtitles(path_replace(episode[0]))
history_log(1, no, episode[2], message)
send_notifications(no, episode[2], message)
list_missing_subtitles(no)
def movies_download_subtitles(no):
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
movie = c_db.execute("SELECT path, missing_subtitles, radarrId, sceneName, hearing_impaired FROM table_movies WHERE radarrId = ?", (no,)).fetchone()
c_db.close()
providers_list = get_providers()
providers_auth = get_providers_auth()
for language in ast.literal_eval(movie[1]):
if language is not None:
message = download_subtitle(path_replace_movie(movie[0]), str(alpha3_from_alpha2(language)), movie[4], providers_list, providers_auth, str(movie[3]), 'movie')
if message is not None:
store_subtitles_movie(path_replace_movie(movie[0]))
history_log_movie(1, no, message)
send_notifications_movie(no, message)
list_missing_subtitles_movies(no)
def wanted_download_subtitles(path):
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
episodes_details = c_db.execute("SELECT table_episodes.path, table_episodes.missing_subtitles, table_episodes.sonarrEpisodeId, table_episodes.sonarrSeriesId, table_shows.hearing_impaired, table_episodes.scene_name, table_episodes.failedAttempts FROM table_episodes INNER JOIN table_shows on table_shows.sonarrSeriesId = table_episodes.sonarrSeriesId WHERE table_episodes.path = ? AND missing_subtitles != '[]'", (path_replace_reverse(path),)).fetchall()
c_db.close()
providers_list = get_providers()
providers_auth = get_providers_auth()
for episode in episodes_details:
attempt = episode[6]
if type(attempt) == unicode:
attempt = ast.literal_eval(attempt)
for language in ast.literal_eval(episode[1]):
if attempt is None:
attempt = []
attempt.append([language, time.time()])
else:
att = zip(*attempt)[0]
if language not in att:
attempt.append([language, time.time()])
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
c_db.execute('UPDATE table_episodes SET failedAttempts = ? WHERE sonarrEpisodeId = ?', (unicode(attempt), episode[2]))
conn_db.commit()
c_db.close()
for i in range(len(attempt)):
if attempt[i][0] == language:
if search_active(attempt[i][1]) is True:
message = download_subtitle(path_replace(episode[0]), str(alpha3_from_alpha2(language)), episode[4], providers_list, providers_auth, str(episode[5]), 'series')
if message is not None:
store_subtitles(path_replace(episode[0]))
list_missing_subtitles(episode[3])
history_log(1, episode[3], episode[2], message)
send_notifications(episode[3], episode[2], message)
else:
logging.debug('Search is not active for episode ' + episode[0] + ' Language: ' + attempt[i][0])
def wanted_download_subtitles_movie(path):
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
movies_details = c_db.execute("SELECT path, missing_subtitles, radarrId, radarrId, hearing_impaired, sceneName, failedAttempts FROM table_movies WHERE path = ? AND missing_subtitles != '[]'", (path_replace_reverse_movie(path),)).fetchall()
c_db.close()
providers_list = get_providers()
providers_auth = get_providers_auth()
for movie in movies_details:
attempt = movie[6]
if type(attempt) == unicode:
attempt = ast.literal_eval(attempt)
for language in ast.literal_eval(movie[1]):
if attempt is None:
attempt = []
attempt.append([language, time.time()])
else:
att = zip(*attempt)[0]
if language not in att:
attempt.append([language, time.time()])
conn_db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
c_db = conn_db.cursor()
c_db.execute('UPDATE table_movies SET failedAttempts = ? WHERE radarrId = ?', (unicode(attempt), movie[2]))
conn_db.commit()
c_db.close()
for i in range(len(attempt)):
if attempt[i][0] == language:
if search_active(attempt[i][1]) is True:
message = download_subtitle(path_replace_movie(movie[0]), str(alpha3_from_alpha2(language)), movie[4], providers_list, providers_auth, str(movie[5]), 'movie')
if message is not None:
store_subtitles_movie(path_replace_movie(movie[0]))
list_missing_subtitles_movies(movie[3])
history_log_movie(1, movie[3], message)
send_notifications_movie(movie[3], message)
else:
logging.info('Search is not active for movie ' + movie[0] + ' Language: ' + attempt[i][0])
def wanted_search_missing_subtitles():
db = sqlite3.connect(os.path.join(config_dir, 'db/bazarr.db'), timeout=30)
db.create_function("path_substitution", 1, path_replace)
db.create_function("path_substitution_movie", 1, path_replace_movie)
c = db.cursor()
if get_general_settings()[24] is True:
monitored_only_query_string = ' AND monitored = "True"'
else:
monitored_only_query_string = ""
c.execute("SELECT path_substitution(path) FROM table_episodes WHERE missing_subtitles != '[]'" + monitored_only_query_string)
episodes = c.fetchall()
c.execute("SELECT path_substitution_movie(path) FROM table_movies WHERE missing_subtitles != '[]'" + monitored_only_query_string)
movies = c.fetchall()
c.close()
integration = get_general_settings()
if integration[12] is True:
for episode in episodes:
wanted_download_subtitles(episode[0])
if integration[13] is True:
for movie in movies:
wanted_download_subtitles_movie(movie[0])
logging.info('Finished searching for missing subtitles. Check histories for more information.')
def search_active(timestamp):
if get_general_settings()[25] is True:
search_deadline = timedelta(weeks=3)
search_delta = timedelta(weeks=1)
aa = datetime.fromtimestamp(float(timestamp))
attempt_datetime = datetime.strptime(str(aa).split(".")[0], '%Y-%m-%d %H:%M:%S')
attempt_search_deadline = attempt_datetime + search_deadline
today = datetime.today()
attempt_age_in_days = (today.date() - attempt_search_deadline.date()).days
if today.date() <= attempt_search_deadline.date():
return True
elif attempt_age_in_days % search_delta.days == 0:
return True
else:
return False
else:
return True