Repository: Diaoul/subliminal
Pull request #1105 (2024-06-03 09:48:25)

Guess subtitle format using pysubs2
Correct #1103  to fix #1104 


Patch:
diff --git a/subliminal/core.py b/subliminal/core.py
index 2c321fe..d3b4ac9 100644
--- a/subliminal/core.py
+++ b/subliminal/core.py
@@ -661,7 +661,7 @@ def download_best_subtitles(videos, languages, min_score=0, hearing_impaired=Fal
     return downloaded_subtitles
 
 
-def save_subtitles(video, subtitles, single=False, directory=None, encoding=None):
+def save_subtitles(video, subtitles, single=False, directory=None, encoding=None, extension=None):
     """Save subtitles on filesystem.
 
     Subtitles are saved in the order of the list. If a subtitle with a language has already been saved, other subtitles
@@ -677,6 +677,7 @@ def save_subtitles(video, subtitles, single=False, directory=None, encoding=None
     :param bool single: save a single subtitle, default is to save one subtitle per language.
     :param str directory: path to directory where to save the subtitles, default is next to the video.
     :param str encoding: encoding in which to save the subtitles, default is to keep original encoding.
+    :param (str | None) extension: the subtitle extension, default is to match to the subtitle format.
     :return: the saved subtitles
     :rtype: list of :class:`~subliminal.subtitle.Subtitle`
 
@@ -694,7 +695,7 @@ def save_subtitles(video, subtitles, single=False, directory=None, encoding=None
             continue
 
         # create subtitle path
-        subtitle_path = subtitle.get_path(video, single=single)
+        subtitle_path = subtitle.get_path(video, single=single, extension=extension)
         if directory is not None:
             subtitle_path = os.path.join(directory, os.path.split(subtitle_path)[1])
 
diff --git a/subliminal/providers/opensubtitlescom.py b/subliminal/providers/opensubtitlescom.py
index 59e4d4f..f3584ff 100644
--- a/subliminal/providers/opensubtitlescom.py
+++ b/subliminal/providers/opensubtitlescom.py
@@ -372,7 +372,7 @@ class OpenSubtitlesComProvider(Provider):
     languages: ClassVar[Set[Language]] = opensubtitlescom_languages
 
     user_agent: str = f'Subliminal v{__short_version__}'
-    expected_format: str = 'srt'
+    subtitle_format: str = 'srt'
 
     username: str | None
     password: str | None
@@ -744,7 +744,7 @@ class OpenSubtitlesComProvider(Provider):
 
         # get the subtitle download link
         logger.info('Downloading subtitle %r', subtitle)
-        body = {'file_id': subtitle.file_id, 'file_name': subtitle.file_name, 'sub_format': self.expected_format}
+        body = {'file_id': subtitle.file_id, 'file_name': subtitle.file_name, 'sub_format': self.subtitle_format}
         r = self.api_post('download', body)
         if any(k not in r for k in ('link', 'remaining')):
             return
diff --git a/subliminal/subtitle.py b/subliminal/subtitle.py
index 081a2b4..c6804a0 100644
--- a/subliminal/subtitle.py
+++ b/subliminal/subtitle.py
@@ -9,7 +9,7 @@ from typing import TYPE_CHECKING, ClassVar
 
 import chardet
 import srt  # type: ignore[import-untyped]
-from pysubs2 import SSAFile  # type: ignore[import-untyped]
+from pysubs2 import SSAFile, UnknownFPSError  # type: ignore[import-untyped]
 
 if TYPE_CHECKING:
     from babelfish import Language  # type: ignore[import-untyped]
@@ -18,8 +18,19 @@ if TYPE_CHECKING:
 
 logger = logging.getLogger(__name__)
 
+#: Subtitle formats to extension
+FORMAT_TO_EXTENSION = {
+    'srt': '.srt',
+    'ass': '.ass',
+    'ssa': '.ssa',
+    'microdvd': '.sub',
+    'mpl2': '.mpl',
+    'tmp': '.txt',
+    'vtt': '.vtt',
+}
+
 #: Subtitle extensions
-SUBTITLE_EXTENSIONS = ('.srt', '.sub', '.smi', '.txt', '.ssa', '.ass', '.mpl', '.vtt')
+SUBTITLE_EXTENSIONS = (*FORMAT_TO_EXTENSION.values(), '.smi')
 
 
 class Subtitle:
@@ -38,9 +49,6 @@ class Subtitle:
     #: Name of the provider that returns that class of subtitle
     provider_name: ClassVar[str] = ''
 
-    #: Subtitle format used to test if the subtitle is valid
-    expected_format: ClassVar[str] = 'srt'
-
     #: Content as bytes
     _content: bytes | None
 
@@ -68,6 +76,9 @@ class Subtitle:
     #: Encoding to decode with when accessing :attr:`text`
     encoding: str | None
 
+    #: Subtitle format, None for automatic detection
+    subtitle_format: str | None = None
+
     #: Framerate for frame-based formats (MicroDVD)
     fps: float | None
 
@@ -78,6 +89,7 @@ class Subtitle:
         hearing_impaired: bool | None = None,
         page_link: str | None = None,
         encoding: str | None = None,
+        subtitle_format: str | None = None,
         fps: float | None = None,
         guess_encoding: bool = True,
     ) -> None:
@@ -90,6 +102,7 @@ class Subtitle:
         self.language = language
         self.hearing_impaired = hearing_impaired
         self.page_link = page_link
+        self.subtitle_format = subtitle_format
         self.fps = fps
 
         self.encoding = None
@@ -194,8 +207,18 @@ class Subtitle:
         if not self.text:
             return False
 
+        # Try guessing the subtitle format
+        if self.subtitle_format is None:
+            guessed_format = get_subtitle_format(self.text, subtitle_format=self.subtitle_format, fps=self.fps)
+            if guessed_format:
+                # Keep the guessed format
+                self.subtitle_format = guessed_format
+                if self.subtitle_format != 'srt':
+                    # Do not check more if the format is not 'srt'
+                    return True
+
         # Valid srt
-        if self.expected_format == 'srt':
+        if self.subtitle_format == 'srt':
             try:
                 parsed = self.parse_srt()
             except Exception:
@@ -206,14 +229,6 @@ class Subtitle:
                     self._text = parsed
                 return True
 
-        # Try other subtitle format
-        try:
-            SSAFile.from_string(self.text, format_=self.expected_format, fps=self.fps)
-        except Exception:
-            logger.exception('not a valid subtitle.')
-        else:
-            return True
-
         return False
 
     def parse_srt(self) -> str:
@@ -334,17 +349,20 @@ class Subtitle:
 
         return encoding_or_none
 
-    def get_path(self, video: Video, *, single: bool = False) -> str:
+    def get_path(self, video: Video, *, single: bool = False, extension: str | None = None) -> str:
         """Get the subtitle path using the `video`, `language` and `extension`.
 
         :param video: path to the video.
         :type video: :class:`~subliminal.video.Video`
         :param bool single: save a single subtitle, default is to save one subtitle per language.
+        :param (str | None) extension: the subtitle extension, default is to match to the subtitle format.
         :return: path of the subtitle.
         :rtype: str
 
         """
-        return get_subtitle_path(video.name, None if single else self.language)
+        if extension is None:
+            extension = FORMAT_TO_EXTENSION.get(self.subtitle_format, '.srt')
+        return get_subtitle_path(video.name, None if single else self.language, extension=extension)
 
     def get_matches(self, video: Video) -> set[str]:
         """Get the matches against the `video`.
@@ -364,6 +382,32 @@ class Subtitle:
         return f'<{self.__class__.__name__} {self.id!r} [{self.language}]>'
 
 
+def get_subtitle_format(
+    text: str,
+    subtitle_format: str | None = None,
+    fps: float | None = None,
+) -> str | None:
+    """Detect the subtitle format with `pysubs2`.
+
+    :param str text: the subtitle text.
+    :param (str | None) subtitle_format: the expected subtitle_format, None for auto-detect.
+    :param (str | None) fps: the framerate for framerate based subtitles.
+    :return: the guessed format or None if not found.
+    :rtype: str | None
+
+    """
+    try:
+        obj = SSAFile.from_string(text, format_=subtitle_format, fps=fps)
+    except UnknownFPSError:
+        default_fps = 24
+        return get_subtitle_format(text, subtitle_format=subtitle_format, fps=default_fps)
+    except Exception:
+        logger.exception('not a valid subtitle.')
+    else:
+        return obj.format
+    return None
+
+
 def get_subtitle_path(video_path: str | os.PathLike, language: Language | None = None, extension: str = '.srt') -> str:
     """Get the subtitle path using the `video_path` and `language`.
 


Test patch:
diff --git a/tests/test_subtitle.py b/tests/test_subtitle.py
index 2c3e91f..c58ec13 100644
--- a/tests/test_subtitle.py
+++ b/tests/test_subtitle.py
@@ -31,6 +31,7 @@ def test_subtitle_is_valid_valid(monkeypatch):
     )
     monkeypatch.setattr(Subtitle, 'text', text)
     assert subtitle.is_valid() is True
+    assert subtitle.subtitle_format == 'srt'
 
 
 @pytest.mark.xfail()
@@ -60,6 +61,22 @@ def test_subtitle_is_valid_valid_begin(monkeypatch):
     assert subtitle.is_valid() is True
 
 
+def test_subtitle_is_valid_sub_format(monkeypatch, movies):
+    video = movies['man_of_steel']
+    subtitle = Subtitle(Language('pol'))
+    text = (
+        '{3146}{3189}/Nie rozumiecie?\n'
+        '{3189}{3244}/Jšdro Kryptona się rozpada.\n'
+        '{3244}{3299}To kwestia tygodni.\n'
+    )
+    monkeypatch.setattr(Subtitle, 'text', text)
+    assert subtitle.is_valid() is True
+    assert subtitle.subtitle_format == 'microdvd'
+    path = subtitle.get_path(video, single=True)
+    extension = os.path.splitext(path)[1]
+    assert extension == '.sub'
+
+
 def test_get_subtitle_path(movies):
     video = movies['man_of_steel']
     assert get_subtitle_path(video.name, extension='.sub') == os.path.splitext(video.name)[0] + '.sub'