1
0
mirror of https://github.com/spaam/svtplay-dl.git synced 2024-11-24 12:15:40 +01:00
svtplay-dl/lib/svtplay_dl/postprocess/__init__.py
Johan Andersson 456e008597 pylint fixes
2017-02-15 23:15:50 +01:00

197 lines
8.3 KiB
Python

from json import dumps
from random import sample
import subprocess
import os
import platform
from requests import post, codes, Timeout
from svtplay_dl.log import log
from svtplay_dl.utils import which, is_py3
class postprocess(object):
def __init__(self, stream, options, subfixes = []):
self.stream = stream
self.merge_subtitle = options.merge_subtitle
self.external_subtitle = options.subtitle
self.get_all_subtitles = options.get_all_subtitles
self.subfixes = subfixes
self.detect = None
for i in ["ffmpeg", "avconv"]:
self.detect = which(i)
if self.detect:
break
def sublanguage(self):
# parse() function partly borrowed from a guy on github. /thanks!
# https://github.com/riobard/srt.py/blob/master/srt.py
def parse(self):
def parse_block(block):
lines = block.strip('-').split('\n')
txt = '\r\n'.join(lines[2:])
return txt
if platform.system() == "Windows" and is_py3:
fd = open(self, encoding="utf8")
else:
fd = open(self)
return list(map(parse_block,
fd.read().strip().replace('\r', '').split('\n\n')))
def query(self):
random_sentences = ' '.join(sample(parse(self), 8)).replace('\r\n', '')
url = 'https://whatlanguage.herokuapp.com'
payload = {"query": random_sentences}
headers = {'content-type': 'application/json'} # Note: requests handles json from version 2.4.2 and onwards so i use json.dumps for now.
try:
r = post(url, data=dumps(payload), headers=headers, timeout=30) # Note: reasonable timeout i guess? svtplay-dl is mainly used while multitasking i presume, and it is heroku after all (fast enough)
if r.status_code == codes.ok:
try:
response = r.json()
return response['language']
except TypeError:
return 'und'
else:
log.error("Server error appeared. Setting language as undetermined.")
return 'und'
except Timeout:
log.error("30 seconds server timeout reached. Setting language as undetermined.")
return 'und'
langs = []
exceptions = {
'lulesamiska': 'smj',
'meankieli': 'fit',
'jiddisch': 'yid'
}
if len(self.subfixes) >= 2:
log.info("Determining the languages of the subtitles.")
else: log.info("Determining the language of the subtitle.")
if self.get_all_subtitles:
from re import match
for subfix in self.subfixes:
if [exceptions[key] for key in exceptions.keys() if match(key, subfix.strip('-'))]:
if 'oversattning' in subfix.strip('-'):
subfix = subfix.strip('-').split('.')[0]
else:
subfix = subfix.strip('-')
langs += [exceptions[subfix]]
continue
subfile = "{0}.srt".format(os.path.splitext(self.stream.options.output)[0] + subfix)
langs += [query(subfile)]
else:
subfile = "{0}.srt".format(os.path.splitext(self.stream.options.output)[0])
langs += [query(subfile)]
if len(langs) >= 2:
log.info("Language codes: " + ', '.join(langs))
else: log.info("Language code: " + langs[0])
return langs
def remux(self):
if self.detect is None:
log.error("Cant detect ffmpeg or avconv. Cant mux files without it.")
return
if self.stream.finished is False:
return
if self.stream.options.output.endswith('.mp4') is False:
orig_filename = self.stream.options.output
name, ext = os.path.splitext(orig_filename)
new_name = u"{0}.mp4".format(name)
if self.merge_subtitle:
log.info(u"Muxing %s and merging its subtitle into %s", orig_filename, new_name)
else:
log.info(u"Muxing %s into %s".format(orig_filename, new_name))
tempfile = u"{0}.temp".format(orig_filename)
arguments = ["-map", "0:v", "-map", "0:a", "-c", "copy", "-copyts", "-f", "mp4"]
if ext == ".ts":
arguments += ["-bsf:a", "aac_adtstoasc"]
cmd = [self.detect, "-i", orig_filename]
if self.merge_subtitle:
langs = self.sublanguage()
for stream_num, language in enumerate(langs):
arguments += ["-map", str(stream_num + 1), "-c:s:" + str(stream_num), "mov_text", "-metadata:s:s:" + str(stream_num), "language=" + language]
if len(self.subfixes) >= 2:
for subfix in self.subfixes:
subfile = "{0}.srt".format(name + subfix)
cmd += ["-i", subfile]
else:
subfile = "{0}.srt".format(name)
cmd += ["-i", subfile]
arguments += ["-y", tempfile]
cmd += arguments
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
stdout, stderr = p.communicate()
if p.returncode != 0:
stderr = stderr.decode('utf-8', 'replace')
msg = stderr.strip().split('\n')[-1]
log.error("Something went wrong: %s", msg)
return
if self.merge_subtitle and not self.external_subtitle:
log.info("Muxing done, removing the old files.")
if len(self.subfixes) >= 2:
for subfix in self.subfixes:
subfile = "{0}.srt".format(name + subfix)
os.remove(subfile)
else: os.remove(subfile)
else: log.info("Muxing done, removing the old file.")
os.remove(orig_filename)
os.rename(tempfile, new_name)
def merge(self):
if self.detect is None:
log.error("Cant detect ffmpeg or avconv. Cant mux files without it.")
return
if self.stream.finished is False:
return
orig_filename = self.stream.options.output
if self.merge_subtitle:
log.info("Merge audio, video and subtitle into %s", orig_filename)
else:
log.info("Merge audio and video into %s", orig_filename)
tempfile = u"{0}.temp".format(orig_filename)
name = os.path.splitext(orig_filename)[0]
audio_filename = u"{0}.m4a".format(name)
arguments = ["-c:v", "copy", "-c:a", "copy", "-f", "mp4"]
cmd = [self.detect, "-i", orig_filename, "-i", audio_filename]
if self.merge_subtitle:
langs = self.sublanguage()
for stream_num, language in enumerate(langs, start=2):
arguments += ["-map", "0", "-map", "1", "-map", str(stream_num), "-c:s:" + str(stream_num - 2), "mov_text", "-metadata:s:s:" + str(stream_num - 2), "language=" + language]
if len(self.subfixes) >= 2:
for subfix in self.subfixes:
subfile = "{0}.srt".format(name + subfix)
cmd += ["-i", subfile]
else:
subfile = "{0}.srt".format(name)
cmd += ["-i", subfile]
arguments += ["-y", tempfile]
cmd += arguments
p = subprocess.Popen(cmd, stdout=subprocess.PIPE, stderr=subprocess.PIPE, stdin=subprocess.PIPE)
stdout, stderr = p.communicate()
if p.returncode != 0:
stderr = stderr.decode('utf-8', 'replace')
msg = stderr.strip().split('\n')[-1]
log.error("Something went wrong: %s", msg)
return
log.info("Merging done, removing old files.")
os.remove(orig_filename)
os.remove(audio_filename)
if self.merge_subtitle and not self.external_subtitle:
if len(self.subfixes) >= 2:
for subfix in self.subfixes:
subfile = "{0}.srt".format(name + subfix)
os.remove(subfile)
else: os.remove(subfile)
os.rename(tempfile, orig_filename)