path: root/libs/pysubs2/ssafile.py

try:
    from collections import abc
except ImportError:
    import collections as abc  # type: ignore[no-redef]
import io
from io import open
from itertools import chain
import os.path
import logging
from typing import Optional, List, Dict, Iterable, Any, overload, Iterator

from .common import IntOrFloat
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):
    """
    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"
    }

    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"``.

    # ------------------------------------------------------------------------
    # I/O methods
    # ------------------------------------------------------------------------

    @classmethod
    def load(cls, path: str, encoding: str="utf-8", format_: Optional[str]=None, fps: Optional[float]=None, **kwargs) -> "SSAFile":
        """
        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()`)

        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.

        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)
            >>> subs3 = pysubs2.load("subrip-subtitles-with-fancy-tags.srt", keep_unknown_html_tags=True)

        """
        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":
        """
        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":
        """
        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):
        """
        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()`)

        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:
        """
        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):
        """
        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):
        """
        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):
        """
        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):
        """
        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):
        """
        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"):
        """
        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:
                    logging.debug("%r missing in self.info", key)
                    return False
                elif other_info is None:
                    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)
                    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

            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:
                    logging.debug("%r missing in self.styles", key)
                    return False
                elif other_style is None:
                    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())
                    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())
                    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)}>"
        else:
            s = f"<SSAFile with 0 events and {len(self.styles)} styles>"

        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:
        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")
        else:
            raise TypeError("Bad key type")

    @overload
    def __delitem__(self, key: int):
        pass

    @overload
    def __delitem__(self, s: slice):
        pass

    def __delitem__(self, key):
        del self.events[key]

    def __len__(self):
        return len(self.events)

    def insert(self, index: int, value: SSAEvent):
        if isinstance(value, SSAEvent):
            self.events.insert(index, value)
        else:
            raise TypeError("SSAFile.events must contain only SSAEvent objects")