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.
bazarr/libs/pysubs2/ssafile.py

521 lines
19 KiB

try:
from collections import abc
except ImportError:
import collections as abc # type: ignore[no-redef]
6 years ago
import io
from io import open
from itertools import chain
6 years ago
import os.path
import logging
from typing import Optional, List, Dict, Iterable, Any, overload, Iterator
from .common import IntOrFloat
6 years ago
from .formats import autodetect_format, get_format_class, get_format_identifier
from .substation import is_valid_field_content
from .ssaevent import SSAEvent
from .ssastyle import SSAStyle
from .time import make_time, ms_to_str
# TODO fix mypy errors regarding SSAFile
class SSAFile(abc.MutableSequence):
6 years ago
"""
Subtitle file in SubStation Alpha format.
This class has a list-like interface which exposes :attr:`SSAFile.events`,
list of subtitles in the file::
subs = SSAFile.load("subtitles.srt")
for line in subs:
print(line.text)
subs.insert(0, SSAEvent(start=0, end=make_time(s=2.5), text="New first subtitle"))
del subs[0]
"""
DEFAULT_INFO = {
"WrapStyle": "0",
"ScaledBorderAndShadow": "yes",
"Collisions": "Normal"
}
6 years ago
def __init__(self):
self.events: List[SSAEvent] = [] #: List of :class:`SSAEvent` instances, ie. individual subtitles.
self.styles: Dict[str, SSAStyle] = {"Default": SSAStyle.DEFAULT_STYLE.copy()} #: Dict of :class:`SSAStyle` instances.
self.info: Dict[str, str] = self.DEFAULT_INFO.copy() #: Dict with script metadata, ie. ``[Script Info]``.
self.aegisub_project: Dict[str, str] = {} #: Dict with Aegisub project, ie. ``[Aegisub Project Garbage]``.
self.fonts_opaque: Dict[str, Any] = {} #: Dict with embedded fonts, ie. ``[Fonts]``.
self.fps: Optional[float] = None #: Framerate used when reading the file, if applicable.
self.format: Optional[str] = None #: Format of source subtitle file, if applicable, eg. ``"srt"``.
6 years ago
# ------------------------------------------------------------------------
# I/O methods
# ------------------------------------------------------------------------
@classmethod
def load(cls, path: str, encoding: str="utf-8", format_: Optional[str]=None, fps: Optional[float]=None, **kwargs) -> "SSAFile":
6 years ago
"""
Load subtitle file from given path.
This method is implemented in terms of :meth:`SSAFile.from_file()`.
See also:
Specific formats may implement additional loading options,
please refer to documentation of the implementation classes
(eg. :meth:`pysubs2.subrip.SubripFormat.from_file()`)
6 years ago
Arguments:
path (str): Path to subtitle file.
encoding (str): Character encoding of input file.
Defaults to UTF-8, you may need to change this.
format_ (str): Optional, forces use of specific parser
(eg. `"srt"`, `"ass"`). Otherwise, format is detected
automatically from file contents. This argument should
be rarely needed.
fps (float): Framerate for frame-based formats (MicroDVD),
for other formats this argument is ignored. Framerate might
be detected from the file, in which case you don't need
to specify it here (when given, this argument overrides
autodetection).
kwargs: Extra options for the reader.
6 years ago
Returns:
SSAFile
Raises:
IOError
UnicodeDecodeError
pysubs2.exceptions.UnknownFPSError
pysubs2.exceptions.UnknownFormatIdentifierError
pysubs2.exceptions.FormatAutodetectionError
Note:
pysubs2 may autodetect subtitle format and/or framerate. These
values are set as :attr:`SSAFile.format` and :attr:`SSAFile.fps`
attributes.
Example:
>>> subs1 = pysubs2.load("subrip-subtitles.srt")
>>> subs2 = pysubs2.load("microdvd-subtitles.sub", fps=23.976)
5 years ago
>>> subs3 = pysubs2.load("subrip-subtitles-with-fancy-tags.srt", keep_unknown_html_tags=True)
6 years ago
"""
with open(path, encoding=encoding) as fp:
return cls.from_file(fp, format_, fps=fps, **kwargs)
@classmethod
def from_string(cls, string: str, format_: Optional[str]=None, fps: Optional[float]=None, **kwargs) -> "SSAFile":
6 years ago
"""
Load subtitle file from string.
See :meth:`SSAFile.load()` for full description.
Arguments:
string (str): Subtitle file in a string. Note that the string
must be Unicode (in Python 2).
Returns:
SSAFile
Example:
>>> text = '''
... 1
... 00:00:00,000 --> 00:00:05,000
... An example SubRip file.
... '''
>>> subs = SSAFile.from_string(text)
"""
fp = io.StringIO(string)
return cls.from_file(fp, format_, fps=fps, **kwargs)
@classmethod
def from_file(cls, fp: io.TextIOBase, format_: Optional[str]=None, fps: Optional[float]=None, **kwargs) -> "SSAFile":
6 years ago
"""
Read subtitle file from file object.
See :meth:`SSAFile.load()` for full description.
Note:
This is a low-level method. Usually, one of :meth:`SSAFile.load()`
or :meth:`SSAFile.from_string()` is preferable.
Arguments:
fp (file object): A file object, ie. :class:`io.TextIOBase` instance.
Note that the file must be opened in text mode (as opposed to binary).
Returns:
SSAFile
"""
if format_ is None:
# Autodetect subtitle format, then read again using correct parser.
# The file might be a pipe and we need to read it twice,
# so just buffer everything.
text = fp.read()
fragment = text[:10000]
format_ = autodetect_format(fragment)
fp = io.StringIO(text)
impl = get_format_class(format_)
subs = cls() # an empty subtitle file
subs.format = format_
subs.fps = fps
impl.from_file(subs, fp, format_, fps=fps, **kwargs)
return subs
def save(self, path: str, encoding: str="utf-8", format_: Optional[str]=None, fps: Optional[float]=None, **kwargs):
6 years ago
"""
Save subtitle file to given path.
This method is implemented in terms of :meth:`SSAFile.to_file()`.
See also:
Specific formats may implement additional saving options,
please refer to documentation of the implementation classes
(eg. :meth:`pysubs2.subrip.SubripFormat.to_file()`)
6 years ago
Arguments:
path (str): Path to subtitle file.
encoding (str): Character encoding of output file.
Defaults to UTF-8, which should be fine for most purposes.
format_ (str): Optional, specifies desired subtitle format
(eg. `"srt"`, `"ass"`). Otherwise, format is detected
automatically from file extension. Thus, this argument
is rarely needed.
fps (float): Framerate for frame-based formats (MicroDVD),
for other formats this argument is ignored. When omitted,
:attr:`SSAFile.fps` value is used (ie. the framerate used
for loading the file, if any). When the :class:`SSAFile`
wasn't loaded from MicroDVD, or if you wish save it with
different framerate, use this argument. See also
:meth:`SSAFile.transform_framerate()` for fixing bad
frame-based to time-based conversions.
kwargs: Extra options for the writer.
Raises:
IOError
UnicodeEncodeError
pysubs2.exceptions.UnknownFPSError
pysubs2.exceptions.UnknownFormatIdentifierError
pysubs2.exceptions.UnknownFileExtensionError
"""
if format_ is None:
ext = os.path.splitext(path)[1].lower()
format_ = get_format_identifier(ext)
with open(path, "w", encoding=encoding) as fp:
self.to_file(fp, format_, fps=fps, **kwargs)
def to_string(self, format_: str, fps: Optional[float]=None, **kwargs) -> str:
6 years ago
"""
Get subtitle file as a string.
See :meth:`SSAFile.save()` for full description.
Returns:
str
"""
fp = io.StringIO()
self.to_file(fp, format_, fps=fps, **kwargs)
return fp.getvalue()
def to_file(self, fp: io.TextIOBase, format_: str, fps: Optional[float]=None, **kwargs):
6 years ago
"""
Write subtitle file to file object.
See :meth:`SSAFile.save()` for full description.
Note:
This is a low-level method. Usually, one of :meth:`SSAFile.save()`
or :meth:`SSAFile.to_string()` is preferable.
Arguments:
fp (file object): A file object, ie. :class:`io.TextIOBase` instance.
Note that the file must be opened in text mode (as opposed to binary).
"""
impl = get_format_class(format_)
impl.to_file(self, fp, format_, fps=fps, **kwargs)
# ------------------------------------------------------------------------
# Retiming subtitles
# ------------------------------------------------------------------------
def shift(self, h: IntOrFloat=0, m: IntOrFloat=0, s: IntOrFloat=0, ms: IntOrFloat=0,
frames: Optional[int]=None, fps: Optional[float]=None):
6 years ago
"""
Shift all subtitles by constant time amount.
Shift may be time-based (the default) or frame-based. In the latter
case, specify both frames and fps. h, m, s, ms will be ignored.
Arguments:
h, m, s, ms: Integer or float values, may be positive or negative.
frames (int): When specified, must be an integer number of frames.
May be positive or negative. fps must be also specified.
fps (float): When specified, must be a positive number.
Raises:
ValueError: Invalid fps or missing number of frames.
"""
delta = make_time(h=h, m=m, s=s, ms=ms, frames=frames, fps=fps)
for line in self:
line.start += delta
line.end += delta
def transform_framerate(self, in_fps: float, out_fps: float):
6 years ago
"""
Rescale all timestamps by ratio of in_fps/out_fps.
Can be used to fix files converted from frame-based to time-based
with wrongly assumed framerate.
Arguments:
in_fps (float)
out_fps (float)
Raises:
ValueError: Non-positive framerate given.
"""
if in_fps <= 0 or out_fps <= 0:
raise ValueError("Framerates must be positive, cannot transform %f -> %f" % (in_fps, out_fps))
ratio = in_fps / out_fps
for line in self:
line.start = int(round(line.start * ratio))
line.end = int(round(line.end * ratio))
# ------------------------------------------------------------------------
# Working with styles
# ------------------------------------------------------------------------
def rename_style(self, old_name: str, new_name: str):
6 years ago
"""
Rename a style, including references to it.
Arguments:
old_name (str): Style to be renamed.
new_name (str): New name for the style (must be unused).
Raises:
KeyError: No style named old_name.
ValueError: new_name is not a legal name (cannot use commas)
or new_name is taken.
"""
if old_name not in self.styles:
raise KeyError("Style %r not found" % old_name)
if new_name in self.styles:
raise ValueError("There is already a style called %r" % new_name)
if not is_valid_field_content(new_name):
raise ValueError("%r is not a valid name" % new_name)
self.styles[new_name] = self.styles[old_name]
del self.styles[old_name]
for line in self:
# XXX also handle \r override tag
if line.style == old_name:
line.style = new_name
def import_styles(self, subs: "SSAFile", overwrite: bool=True):
6 years ago
"""
Merge in styles from other SSAFile.
Arguments:
subs (SSAFile): Subtitle file imported from.
overwrite (bool): On name conflict, use style from the other file
(default: True).
"""
if not isinstance(subs, SSAFile):
raise TypeError("Must supply an SSAFile.")
for name, style in subs.styles.items():
if name not in self.styles or overwrite:
self.styles[name] = style
# ------------------------------------------------------------------------
# Helper methods
# ------------------------------------------------------------------------
def remove_miscellaneous_events(self):
"""
Remove subtitles which appear to be non-essential (the --clean in CLI)
Currently, this removes events matching any of these criteria:
- SSA event type Comment
- SSA drawing tags
- Less than two characters of text
- Duplicated text with identical time interval (only the first event is kept)
"""
new_events = []
duplicate_text_ids = set()
times_to_texts = {}
for i, e in enumerate(self):
tmp = times_to_texts.setdefault((e.start, e.end), [])
if tmp.count(e.plaintext) > 0:
duplicate_text_ids.add(i)
tmp.append(e.plaintext)
for i, e in enumerate(self):
if e.is_drawing or e.is_comment:
continue
if len(e.plaintext.strip()) < 2:
continue
if i in duplicate_text_ids:
continue
new_events.append(e)
self.events = new_events
def equals(self, other: "SSAFile"):
6 years ago
"""
Equality of two SSAFiles.
Compares :attr:`SSAFile.info`, :attr:`SSAFile.styles` and :attr:`SSAFile.events`.
Order of entries in OrderedDicts does not matter. "ScriptType" key in info is
considered an implementation detail and thus ignored.
Useful mostly in unit tests. Differences are logged at DEBUG level.
"""
if isinstance(other, SSAFile):
for key in set(chain(self.info.keys(), other.info.keys())) - {"ScriptType"}:
self_info, other_info = self.info.get(key), other.info.get(key)
if self_info is None:
6 years ago
logging.debug("%r missing in self.info", key)
return False
elif other_info is None:
6 years ago
logging.debug("%r missing in other.info", key)
return False
elif self_info != other_info:
logging.debug("info %r differs (self=%r, other=%r)", key, self_info, other_info)
6 years ago
return False
for key in set(chain(self.fonts_opaque.keys(), other.fonts_opaque.keys())):
self_font, other_font = self.fonts_opaque.get(key), other.fonts_opaque.get(key)
if self_font is None:
logging.debug("%r missing in self.fonts_opaque", key)
return False
elif other_font is None:
logging.debug("%r missing in other.fonts_opaque", key)
return False
elif self_font != other_font:
logging.debug("fonts_opaque %r differs (self=%r, other=%r)", key, self_font, other_font)
return False
6 years ago
for key in set(chain(self.styles.keys(), other.styles.keys())):
self_style, other_style = self.styles.get(key), other.styles.get(key)
if self_style is None:
6 years ago
logging.debug("%r missing in self.styles", key)
return False
elif other_style is None:
6 years ago
logging.debug("%r missing in other.styles", key)
return False
elif self_style != other_style:
for k in self_style.FIELDS:
if getattr(self_style, k) != getattr(other_style, k): logging.debug("difference in field %r", k)
logging.debug("style %r differs (self=%r, other=%r)", key, self_style.as_dict(), other_style.as_dict())
6 years ago
return False
if len(self) != len(other):
logging.debug("different # of subtitles (self=%d, other=%d)", len(self), len(other))
return False
for i, (self_event, other_event) in enumerate(zip(self.events, other.events)):
if not self_event.equals(other_event):
for k in self_event.FIELDS:
if getattr(self_event, k) != getattr(other_event, k): logging.debug("difference in field %r", k)
logging.debug("event %d differs (self=%r, other=%r)", i, self_event.as_dict(), other_event.as_dict())
6 years ago
return False
return True
else:
raise TypeError("Cannot compare to non-SSAFile object")
def __repr__(self):
if self.events:
max_time = max(ev.end for ev in self)
s = f"<SSAFile with {len(self)} events and {len(self.styles)} styles, last timestamp {ms_to_str(max_time)}>"
6 years ago
else:
s = f"<SSAFile with 0 events and {len(self.styles)} styles>"
6 years ago
return s
# ------------------------------------------------------------------------
# MutableSequence implementation + sort()
# ------------------------------------------------------------------------
def sort(self):
"""Sort subtitles time-wise, in-place."""
self.events.sort()
def __iter__(self) -> Iterator[SSAEvent]:
return iter(self.events)
@overload
def __getitem__(self, item: int) -> SSAEvent:
6 years ago
return self.events[item]
@overload
def __getitem__(self, s: slice) -> List[SSAEvent]:
return self.events[s]
def __getitem__(self, item):
return self.events[item]
@overload
def __setitem__(self, key: int, value: SSAEvent):
pass
@overload
def __setitem__(self, keys: slice, values: Iterable[SSAEvent]):
pass
def __setitem__(self, key, value):
if isinstance(key, int):
if isinstance(value, SSAEvent):
self.events[key] = value
else:
raise TypeError("SSAFile.events must contain only SSAEvent objects")
elif isinstance(key, slice):
values = list(value)
if all(isinstance(v, SSAEvent) for v in values):
self.events[key] = values
else:
raise TypeError("SSAFile.events must contain only SSAEvent objects")
6 years ago
else:
raise TypeError("Bad key type")
6 years ago
@overload
def __delitem__(self, key: int):
pass
@overload
def __delitem__(self, s: slice):
pass
def __delitem__(self, key):
6 years ago
del self.events[key]
def __len__(self):
return len(self.events)
def insert(self, index: int, value: SSAEvent):
6 years ago
if isinstance(value, SSAEvent):
self.events.insert(index, value)
else:
raise TypeError("SSAFile.events must contain only SSAEvent objects")