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()