File size: 2,059 Bytes
969573c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 | from __future__ import annotations
import json
import unittest
from moss_transcribe_diarize.subtitle import SubtitleSegment, SubtitleStyle, export_ass, export_json, export_srt
from moss_transcribe_diarize.subtitle.export import format_ass_time, format_srt_time
class SubtitleExportTest(unittest.TestCase):
def test_time_formatters(self):
self.assertEqual(format_srt_time(3661.234), "01:01:01,234")
self.assertEqual(format_ass_time(3661.23), "1:01:01.23")
def test_export_srt(self):
text = export_srt([SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")])
self.assertIn("00:00:00,500 --> 00:00:02,000", text)
self.assertIn("S01: hello", text)
def test_export_srt_with_speaker_names(self):
text = export_srt(
[SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")],
speaker_names={"S01": "Alice"},
)
self.assertIn("Alice: hello", text)
def test_export_ass(self):
text = export_ass(
[SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")],
style=SubtitleStyle(font_size=42, show_speaker=False),
video_width=1280,
video_height=720,
)
self.assertIn("PlayResX: 1280", text)
self.assertIn("Style: Speaker_S01,Noto Sans CJK SC,42", text)
self.assertIn("Dialogue: 0,0:00:00.50,0:00:02.00,Speaker_S01", text)
self.assertIn("hello", text)
def test_export_ass_with_speaker_names(self):
text = export_ass(
[SubtitleSegment("seg_0001", 0.5, 2.0, "S01", "hello")],
style=SubtitleStyle(font_size=42, speaker_names={"S01": "Alice"}),
video_width=1280,
video_height=720,
)
self.assertIn("Alice: hello", text)
def test_export_json(self):
data = json.loads(export_json([SubtitleSegment("seg_0001", 0, 1, "S01", "hello")]))
self.assertEqual(data[0]["id"], "seg_0001")
self.assertEqual(data[0]["text"], "hello")
if __name__ == "__main__":
unittest.main()
|