diff --git a/tests/test_vtt.py b/tests/test_vtt.py new file mode 100644 index 000000000..a9099f7bb --- /dev/null +++ b/tests/test_vtt.py @@ -0,0 +1,79 @@ +from copy import deepcopy +from io import StringIO + +import pytest + +from whisper.utils import WriteSRT, WriteVTT + + +@pytest.mark.parametrize( + "text, expected", + [ + ("", "<name>"), + ("< & &", "&lt; &amp; &"), + ("name", "<u>name</u>"), + ("a > b", "a > b"), + ('"quoted" text', '"quoted" text'), + ("hello 世界", "hello 世界"), + ], +) +def test_vtt_escapes_segment_text(text, expected): + result = {"segments": [{"start": 0.0, "end": 1.0, "text": text}]} + original = deepcopy(result) + output = StringIO() + WriteVTT(".").write_result(result, output) + assert output.getvalue() == f"WEBVTT\n\n00:00.000 --> 00:01.000\n{expected}\n\n" + assert result == original + + +@pytest.mark.parametrize("highlight", [False, True]) +def test_vtt_escapes_words_before_highlighting(highlight): + words = [ + {"word": " name", "start": 0.0, "end": 0.5}, + {"word": " <", "start": 1.0, "end": 1.5}, + ] + result = {"segments": [{"start": 0.0, "end": 1.5, "words": words}]} + original = deepcopy(result) + cues = list(WriteVTT(".").iterate_result(result, highlight_words=highlight)) + text = "<u>name</u> &lt;" + expected = ( + [ + ("00:00.000", "00:00.500", "<u>name</u> &lt;"), + ("00:00.500", "00:01.000", text), + ("00:01.000", "00:01.500", "<u>name</u> &lt;"), + ] + if highlight + else [("00:00.000", "00:01.500", text)] + ) + assert cues == expected + assert result == original + + +@pytest.mark.parametrize( + "writer_class, expected", + [(WriteVTT, "<x> &\ny"), (WriteSRT, " &\ny")], +) +def test_subtitle_wrapping_uses_unescaped_text(writer_class, expected): + words = [ + {"word": " ", "start": 0.0, "end": 1.0}, + {"word": " &", "start": 1.0, "end": 2.0}, + {"word": " y", "start": 2.0, "end": 3.0}, + ] + result = {"segments": [{"start": 0.0, "end": 3.0, "words": words}]} + original = deepcopy(result) + cues = list(writer_class(".").iterate_result(result, max_line_width=5)) + assert len(cues) == 1 + assert cues[0][2] == expected + assert result == original + + +@pytest.mark.parametrize("with_words", [False, True]) +def test_srt_keeps_existing_text_format(with_words): + text = "name <" + segment = {"start": 0.0, "end": 1.0, "text": text} + if with_words: + segment["words"] = [{"word": text, "start": 0.0, "end": 1.0}] + result = {"segments": [segment]} + assert list(WriteSRT(".").iterate_result(result)) == [ + ("00:00:00,000", "00:00:01,000", text) + ] diff --git a/whisper/utils.py b/whisper/utils.py index 79658fc75..aaaa505d7 100644 --- a/whisper/utils.py +++ b/whisper/utils.py @@ -1,3 +1,4 @@ +import html import json import os import re @@ -197,10 +198,10 @@ def iterate_subtitles(): for subtitle in iterate_subtitles(): subtitle_start = self.format_timestamp(subtitle[0]["start"]) subtitle_end = self.format_timestamp(subtitle[-1]["end"]) - subtitle_text = "".join([word["word"] for word in subtitle]) + all_words = [self._escape_text(timing["word"]) for timing in subtitle] + subtitle_text = "".join(all_words) if highlight_words: last = subtitle_start - all_words = [timing["word"] for timing in subtitle] for i, this_word in enumerate(subtitle): start = self.format_timestamp(this_word["start"]) end = self.format_timestamp(this_word["end"]) @@ -225,7 +226,10 @@ def iterate_subtitles(): segment_start = self.format_timestamp(segment["start"]) segment_end = self.format_timestamp(segment["end"]) segment_text = segment["text"].strip().replace("-->", "->") - yield segment_start, segment_end, segment_text + yield segment_start, segment_end, self._escape_text(segment_text) + + def _escape_text(self, text: str) -> str: + return text def format_timestamp(self, seconds: float): return format_timestamp( @@ -240,6 +244,9 @@ class WriteVTT(SubtitlesWriter): always_include_hours: bool = False decimal_marker: str = "." + def _escape_text(self, text: str) -> str: + return html.escape(text, quote=False) + def write_result( self, result: dict, file: TextIO, options: Optional[dict] = None, **kwargs ):