You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
164 lines
4.9 KiB
164 lines
4.9 KiB
# -*- coding: utf-8 -*-
|
|
import copy
|
|
from datetime import timedelta
|
|
import logging
|
|
import os
|
|
|
|
import pysubs2
|
|
import srt
|
|
import six
|
|
import sys
|
|
|
|
logging.basicConfig(level=logging.INFO)
|
|
logger = logging.getLogger(__name__)
|
|
|
|
|
|
class SubsMixin(object):
|
|
def __init__(self, subs=None):
|
|
self.subs_ = subs
|
|
|
|
def set_encoding(self, encoding):
|
|
self.subs_.set_encoding(encoding)
|
|
return self
|
|
|
|
|
|
class GenericSubtitle(object):
|
|
def __init__(self, start, end, inner):
|
|
self.start = start
|
|
self.end = end
|
|
self.inner = inner
|
|
|
|
def __eq__(self, other):
|
|
eq = True
|
|
eq = eq and self.start == other.start
|
|
eq = eq and self.end == other.end
|
|
eq = eq and self.inner == other.inner
|
|
return eq
|
|
|
|
@property
|
|
def content(self):
|
|
if isinstance(self.inner, srt.Subtitle):
|
|
ret = self.inner.content
|
|
elif isinstance(self.inner, pysubs2.SSAEvent):
|
|
ret = self.inner.text
|
|
else:
|
|
raise NotImplementedError('unsupported subtitle type: %s' % type(self.inner))
|
|
return ret
|
|
|
|
def resolve_inner_timestamps(self):
|
|
ret = copy.deepcopy(self.inner)
|
|
if isinstance(self.inner, srt.Subtitle):
|
|
ret.start = self.start
|
|
ret.end = self.end
|
|
elif isinstance(self.inner, pysubs2.SSAEvent):
|
|
ret.start = pysubs2.make_time(s=self.start.total_seconds())
|
|
ret.end = pysubs2.make_time(s=self.end.total_seconds())
|
|
else:
|
|
raise NotImplementedError('unsupported subtitle type: %s' % type(self.inner))
|
|
return ret
|
|
|
|
def merge_with(self, other):
|
|
assert isinstance(self.inner, type(other.inner))
|
|
inner_merged = copy.deepcopy(self.inner)
|
|
if isinstance(self.inner, srt.Subtitle):
|
|
inner_merged.content = u'{}\n{}'.format(inner_merged.content, other.inner.content)
|
|
return self.__class__(
|
|
self.start,
|
|
self.end,
|
|
inner_merged
|
|
)
|
|
else:
|
|
raise NotImplementedError('unsupported subtitle type: %s' % type(self.inner))
|
|
|
|
@classmethod
|
|
def wrap_inner_subtitle(cls, sub):
|
|
if isinstance(sub, srt.Subtitle):
|
|
return cls(sub.start, sub.end, sub)
|
|
elif isinstance(sub, pysubs2.SSAEvent):
|
|
return cls(
|
|
timedelta(milliseconds=sub.start),
|
|
timedelta(milliseconds=sub.end),
|
|
sub
|
|
)
|
|
else:
|
|
raise NotImplementedError('unsupported subtitle type: %s' % type(sub))
|
|
|
|
|
|
class GenericSubtitlesFile(object):
|
|
def __init__(self, subs, *args, **kwargs):
|
|
sub_format = kwargs.pop('sub_format', None)
|
|
if sub_format is None:
|
|
raise ValueError('format must be specified')
|
|
encoding = kwargs.pop('encoding', None)
|
|
if encoding is None:
|
|
raise ValueError('encoding must be specified')
|
|
self.subs_ = subs
|
|
self._sub_format = sub_format
|
|
self._encoding = encoding
|
|
self._styles = kwargs.pop('styles', None)
|
|
|
|
def set_encoding(self, encoding):
|
|
if encoding != 'same':
|
|
self._encoding = encoding
|
|
return self
|
|
|
|
def __len__(self):
|
|
return len(self.subs_)
|
|
|
|
def __getitem__(self, item):
|
|
return self.subs_[item]
|
|
|
|
@property
|
|
def sub_format(self):
|
|
return self._sub_format
|
|
|
|
@property
|
|
def encoding(self):
|
|
return self._encoding
|
|
|
|
@property
|
|
def styles(self):
|
|
return self._styles
|
|
|
|
def gen_raw_resolved_subs(self):
|
|
for sub in self.subs_:
|
|
yield sub.resolve_inner_timestamps()
|
|
|
|
def offset(self, td):
|
|
offset_subs = []
|
|
for sub in self.subs_:
|
|
offset_subs.append(
|
|
GenericSubtitle(sub.start + td, sub.end + td, sub.inner)
|
|
)
|
|
return GenericSubtitlesFile(
|
|
offset_subs,
|
|
sub_format=self.sub_format,
|
|
encoding=self.encoding,
|
|
styles=self.styles
|
|
)
|
|
|
|
def write_file(self, fname):
|
|
# TODO: converter to go between self.subs_format and out_format
|
|
if fname is None:
|
|
out_format = self._sub_format
|
|
else:
|
|
out_format = os.path.splitext(fname)[-1][1:]
|
|
subs = list(self.gen_raw_resolved_subs())
|
|
if out_format == 'srt':
|
|
to_write = srt.compose(subs)
|
|
elif out_format in ('ssa', 'ass'):
|
|
ssaf = pysubs2.SSAFile()
|
|
ssaf.events = subs
|
|
ssaf.styles = self.styles
|
|
to_write = ssaf.to_string(out_format)
|
|
else:
|
|
raise NotImplementedError('unsupported output format: %s' % out_format)
|
|
|
|
to_write = to_write.encode(self.encoding)
|
|
if six.PY3:
|
|
with open(fname or sys.stdout.fileno(), 'wb') as f:
|
|
f.write(to_write)
|
|
else:
|
|
with (fname and open(fname, 'wb')) or sys.stdout as f:
|
|
f.write(to_write)
|