(VTK-2583) Support outputting VTT with sequence numbers

pycaption/webvtt.py

@@ -197,7 +197,7 @@ class WebVTTWriter(BaseWriter):
     video_width = None
     video_height = None
 
-    def write(self, caption_set, lang=None, force_hours=False):
+    def write(self, caption_set, lang=None, force_hours=False, include_sequence_numbers=False):
         """
         :type caption_set: CaptionSet
         @param force_hours: force writing timestamps in full (hh:mm:ss.xxx) even when "hour" is 0
@@ -220,10 +220,15 @@ def write(self, caption_set, lang=None, force_hours=False):
         self.global_layout = caption_set.get_layout_info(lang)
 
         captions = caption_set.get_captions(lang)
+        captions_out = []
+        count = 1 if include_sequence_numbers else None
+        for caption in captions:
+            converted_caption = self._convert_caption(caption_set, caption, force_hours, count=count)
+            captions_out.append(converted_caption)
+            if count is not None:
+                count += 1
 
-        return output + "\n".join(
-            [self._convert_caption(caption_set, caption, force_hours) for caption in captions]
-        )
+        return output + "\n".join(captions_out)
 
     def _timestamp(self, ts, force_hours):
         td = datetime.timedelta(microseconds=ts)
@@ -265,7 +270,7 @@ def _calculate_resulting_style(self, style, caption_set):
 
         return resulting_style
 
-    def _convert_caption(self, caption_set, caption, force_hours):
+    def _convert_caption(self, caption_set, caption, force_hours, count=None):
         """
         :type caption: Caption
         """
@@ -291,6 +296,8 @@ def _convert_caption(self, caption_set, caption, force_hours):
             if not layout:
                 layout = caption.layout_info or self.global_layout
             cue_settings = self._convert_positioning(layout)
+            if count is not None:
+                output += f"{count}\n"
             output += timespan + cue_settings + "\n"
             output += cue_style_tags[0] + cue_text + cue_style_tags[1] + "\n"
 

tests/conftest.py

@@ -94,7 +94,7 @@
     sample_webvtt_from_dfxp_with_style,
     sample_webvtt_keeps_positioning,
     sample_webvtt_from_dfxp_with_positioning_and_style,
-    sample_webvtt_from_srt, sample_webvtt_from_webvtt,
+    sample_webvtt_from_srt, sample_webvtt_from_srt_with_sequence, sample_webvtt_from_webvtt,
     sample_webvtt_2, sample_webvtt_empty, sample_webvtt_double_br,
     sample_webvtt_output_long_cue, webvtt_from_dfxp_with_conflicting_align,
     sample_webvtt_with_cue_settings,

tests/fixtures/webvtt.py

@@ -245,6 +245,45 @@ def sample_webvtt_from_srt():
 <LAUGHING & WHOOPS!>
 """
 
+@pytest.fixture(scope="session")
+def sample_webvtt_from_srt_with_sequence():
+    return """WEBVTT
+
+1
+00:09.209 --> 00:12.312
+( clock ticking )
+
+2
+00:14.848 --> 00:17.000
+MAN:
+When we think
+♪ ...say bow, wow, ♪
+
+3
+00:17.000 --> 00:18.752
+we have this vision of Einstein
+
+4
+00:18.752 --> 00:20.887
+as an old, wrinkly man
+with white hair.
+
+5
+00:20.887 --> 00:26.760
+MAN 2:
+E equals m c-squared is
+not about an old Einstein.
+
+6
+00:26.760 --> 00:32.200
+MAN 2:
+It's all about an eternal Einstein.
+
+7
+00:32.200 --> 00:36.200
+<LAUGHING & WHOOPS!>
+"""
+
 
 # This is not equal to the input because we accept unescaped illegal characters
 # when reading (because many players do so) but escape them when writing

tests/test_webvtt_conversion.py

@@ -59,6 +59,15 @@ def test_srt_to_webvtt_conversion(self, sample_webvtt_from_srt, sample_srt):
         self.assert_webvtt_equals(sample_webvtt_from_srt, results)
 
 
+class TestSRTtoWebVTTWithSequence(WebVTTTestingMixIn):
+    def test_srt_to_webvtt_with_sequence_conversion(self, sample_webvtt_from_srt_with_sequence, sample_srt):
+        caption_set = SRTReader().read(sample_srt)
+        results = WebVTTWriter().write(caption_set, include_sequence_numbers=True)
+
+        assert isinstance(results, str)
+        self.assert_webvtt_equals(sample_webvtt_from_srt_with_sequence, results)
+
+
 class TestWebVTTtoWebVTT(WebVTTTestingMixIn):
     def test_webvtt_to_webvtt_conversion(self, sample_webvtt_from_webvtt,
                                          sample_webvtt):