feat: support create subtitles (.srt)

test_subtitles.py

@@ -0,0 +1,75 @@
+import unittest
+
+from subtitles import convert, timestamp
+
+
+class SubtitleTests(unittest.TestCase):
+    def test_minute_long_segment_becomes_short_cues(self):
+        text = (
+            "Okay, all right, brilliant. Oh dude, we've got it this time. "
+            "You know what? I can feel it. Okay, we found the longer self "
+            "tappers that are for these little SSD cooling fans. "
+        ) * 4
+        tokens = text.split()
+        words = [
+            {"word": word, "start": 971.68 + i * 0.4, "end": 971.68 + i * 0.4 + 0.35}
+            for i, word in enumerate(tokens)
+        ]
+        result = convert(
+            {
+                "text": text,
+                "words": words,
+                "segments": [{"text": text, "start": 971.68, "end": 1037.42}],
+            }
+        )
+        cues = result.strip().split("\n\n")
+        self.assertGreater(len(cues), 10)
+        recovered = []
+        for index, cue in enumerate(cues, 1):
+            number, timing, *body = cue.splitlines()
+            self.assertEqual(int(number), index)
+            self.assertLessEqual(len(body), 2)
+            self.assertTrue(all(len(line) <= 42 for line in body))
+
+            def seconds(value):
+                h, m, s = value.replace(",", ".").split(":")
+                return int(h) * 3600 + int(m) * 60 + float(s)
+
+            start, end = map(seconds, timing.split(" --> "))
+            self.assertLessEqual(end - start, 6.001)
+            recovered.extend(" ".join(body).split())
+        self.assertEqual(recovered, tokens)
+        self.assertIn("00:16:11,680", result)
+
+    def test_pause_sentence_and_speaker_boundaries(self):
+        result = convert(
+            {
+                "words": [
+                    {"word": "First sentence.", "start": 0, "end": 1.2, "speaker": 0},
+                    {"word": "Next", "start": 1.3, "end": 1.6, "speaker": 0},
+                    {"word": "pause", "start": 3, "end": 3.4, "speaker": 0},
+                    {"word": "speaker", "start": 3.5, "end": 4, "speaker": 1},
+                ]
+            }
+        )
+        self.assertEqual(len(result.strip().split("\n\n")), 4)
+
+    def test_missing_and_invalid_timestamps(self):
+        for response in (
+            {"text": "speech"},
+            {"text": "speech", "words": []},
+            {"words": [{"word": "bad", "start": 2, "end": 1}]},
+            {"words": [{"word": "bad", "start": float("nan"), "end": 1}]},
+            {"error": {"message": "failed"}},
+        ):
+            with self.assertRaises(ValueError):
+                convert(response)
+
+    def test_empty_and_rounding(self):
+        self.assertEqual(convert({"words": [], "text": ""}), "")
+        self.assertEqual(timestamp(59.9996), "00:01:00,000")
+        self.assertEqual(timestamp(3601.234), "01:00:01,234")
+
+
+if __name__ == "__main__":
+    unittest.main()