Add support for on-demand subtitles (#1305)

JonnyWong16 · web-flow · commit d4426dab86fc · 2023-12-22T13:13:04.000-08:00
* Add support for on-demand subtitle search * Clean up uploadSubtitles and removeSubtitles methods * Add test for on-demand subtitles Account is required for on-demand subtitles * Add hearingImpaired and perfectMatch to SubtitleStream attributes * Update subtitle score doc string The score is the OpenSubtitles download count Ref: https://forums.plex.tv/t/subtitle-search-update-opensubtitles/862746
diff --git a/plexapi/media.py b/plexapi/media.py
@@ -444,8 +444,10 @@ class SubtitleStream(MediaPartStream):
             forced (bool): True if this is a forced subtitle.
             format (str): The format of the subtitle stream (ex: srt).
             headerCompression (str): The header compression of the subtitle stream.
+            hearingImpaired (bool): True if this is a hearing impaired (SDH) subtitle.
+            perfectMatch (bool): True if the on-demand subtitle is a perfect match.
             providerTitle (str): The provider title where the on-demand subtitle is downloaded from.
-            score (int): The match score of the on-demand subtitle.
+            score (int): The match score (download count) of the on-demand subtitle.
             sourceKey (str): The source key of the on-demand subtitle.
             transient (str): Unknown.
             userID (int): The user id of the user that downloaded the on-demand subtitle.
@@ -460,6 +462,8 @@ def _loadData(self, data):
         self.forced = utils.cast(bool, data.attrib.get('forced', '0'))
         self.format = data.attrib.get('format')
         self.headerCompression = data.attrib.get('headerCompression')
+        self.hearingImpaired = utils.cast(bool, data.attrib.get('hearingImpaired', '0'))
+        self.perfectMatch = utils.cast(bool, data.attrib.get('perfectMatch'))
         self.providerTitle = data.attrib.get('providerTitle')
         self.score = utils.cast(int, data.attrib.get('score'))
         self.sourceKey = data.attrib.get('sourceKey')
diff --git a/plexapi/video.py b/plexapi/video.py
@@ -98,27 +98,87 @@ def _defaultSyncTitle(self):
         return self.title
 
     def uploadSubtitles(self, filepath):
-        """ Upload Subtitle file for video. """
+        """ Upload a subtitle file for the video.
+
+            Parameters:
+                filepath (str): Path to subtitle file.
+        """
         url = f'{self.key}/subtitles'
         filename = os.path.basename(filepath)
         subFormat = os.path.splitext(filepath)[1][1:]
+        params = {
+            'title': filename,
+            'format': subFormat,
+        }
+        headers = {'Accept': 'text/plain, */*'}
         with open(filepath, 'rb') as subfile:
-            params = {'title': filename,
-                      'format': subFormat
-                      }
-            headers = {'Accept': 'text/plain, */*'}
             self._server.query(url, self._server._session.post, data=subfile, params=params, headers=headers)
         return self
 
-    def removeSubtitles(self, streamID=None, streamTitle=None):
-        """ Remove Subtitle from movie's subtitles listing.
+    def searchSubtitles(self, language='en', hearingImpaired=0, forced=0):
+        """ Search for on-demand subtitles for the video.
+            See https://support.plex.tv/articles/subtitle-search/.
 
-            Note: If subtitle file is located inside video directory it will bbe deleted.
-            Files outside of video directory are not effected.
+            Parameters:
+                language (str, optional): Language code (ISO 639-1) of the subtitles to search for.
+                    Default 'en'.
+                hearingImpaired (int, optional): Search option for SDH subtitles.
+                    Default 0.
+                    (0 = Prefer non-SDH subtitles, 1 = Prefer SDH subtitles,
+                    2 = Only show SDH subtitles, 3 = Only show non-SDH subtitles)
+                forced (int, optional): Search option for forced subtitles.
+                    Default 0.
+                    (0 = Prefer non-forced subtitles, 1 = Prefer forced subtitles,
+                    2 = Only show forced subtitles, 3 = Only show non-forced subtitles)
+
+            Returns:
+                List<:class:`~plexapi.media.SubtitleStream`>: List of SubtitleStream objects.
+        """
+        params = {
+            'language': language,
+            'hearingImpaired': hearingImpaired,
+            'forced': forced,
+        }
+        key = f'{self.key}/subtitles{utils.joinArgs(params)}'
+        return self.fetchItems(key)
+
+    def downloadSubtitles(self, subtitleStream):
+        """ Download on-demand subtitles for the video.
+            See https://support.plex.tv/articles/subtitle-search/.
+
+            Note: This method is asynchronous and returns immediately before subtitles are fully downloaded.
+
+            Parameters:
+                subtitleStream (:class:`~plexapi.media.SubtitleStream`):
+                    Subtitle object returned from :func:`~plexapi.video.Video.searchSubtitles`.
+        """
+        key = f'{self.key}/subtitles'
+        params = {'key': subtitleStream.key}
+        self._server.query(key, self._server._session.put, params=params)
+        return self
+
+    def removeSubtitles(self, subtitleStream=None, streamID=None, streamTitle=None):
+        """ Remove an upload or downloaded subtitle from the video.
+
+            Note: If the subtitle file is located inside video directory it will be deleted.
+            Files outside of video directory are not affected.
+            Embedded subtitles cannot be removed.
+
+            Parameters:
+                subtitleStream (:class:`~plexapi.media.SubtitleStream`, optional): Subtitle object to remove.
+                streamID (int, optional): ID of the subtitle stream to remove.
+                streamTitle (str, optional): Title of the subtitle stream to remove.
         """
-        for stream in self.subtitleStreams():
-            if streamID == stream.id or streamTitle == stream.title:
-                self._server.query(stream.key, self._server._session.delete)
+        if subtitleStream is None:
+            try:
+                subtitleStream = next(
+                    stream for stream in self.subtitleStreams()
+                    if streamID == stream.id or streamTitle == stream.title
+                )
+            except StopIteration:
+                raise BadRequest(f"Subtitle stream with ID '{streamID}' or title '{streamTitle}' not found.") from None
+
+        self._server.query(subtitleStream.key, self._server._session.delete)
         return self
 
     def optimize(self, title='', target='', deviceProfile='', videoQuality=None,
diff --git a/tests/test_video.py b/tests/test_video.py
@@ -404,7 +404,6 @@ def test_video_Episode_subtitleStreams(episode):
 
 
 def test_video_Movie_upload_select_remove_subtitle(movie, subtitle):
-
     filepath = os.path.realpath(subtitle.name)
 
     movie.uploadSubtitles(filepath)
@@ -414,7 +413,6 @@ def test_video_Movie_upload_select_remove_subtitle(movie, subtitle):
 
     movie.subtitleStreams()[0].setSelected()
     movie.reload()
-
     subtitleSelection = movie.subtitleStreams()[0]
     assert subtitleSelection.selected
 
@@ -429,6 +427,22 @@ def test_video_Movie_upload_select_remove_subtitle(movie, subtitle):
         pass
 
 
+def test_video_Movie_on_demand_subtitles(movie, account):
+    movie_subtitles = movie.subtitleStreams()
+    subtitles = movie.searchSubtitles()
+    assert subtitles != []
+
+    subtitle = subtitles[0]
+
+    movie.downloadSubtitles(subtitle)
+    utils.wait_until(lambda: len(movie.reload().subtitleStreams()) > len(movie_subtitles))
+    subtitle_sourceKeys = {stream.sourceKey: stream for stream in movie.subtitleStreams()}
+    assert subtitle.sourceKey in subtitle_sourceKeys
+
+    movie.removeSubtitles(subtitleStream=subtitle_sourceKeys[subtitle.sourceKey]).reload()
+    assert subtitle.sourceKey not in [stream.sourceKey for stream in movie.subtitleStreams()]
+
+
 def test_video_Movie_match(movies):
     sectionAgent = movies.agent
     sectionAgents = [agent.identifier for agent in movies.agents() if agent.shortIdentifier != 'none']
