From 0d13013330986d0989db024b25f9b382473d2144 Mon Sep 17 00:00:00 2001 From: Seongho Bae Date: Mon, 6 Jul 2026 08:59:41 +0900 Subject: [PATCH] feat: add subtitle/caption export module (.srt and .vtt) Add a self-contained, stdlib-only subtitles.py that converts timestamped transcript segments into SubRip (.srt) and WebVTT (.vtt) captions. - Segment dataclass as a tiny input contract (also accepts duck-typed objects with .start/.end/.text), independent of transcribe.py - to_srt / to_vtt render formatted text; write_srt / write_vtt write files - Correct timestamp formatting: SRT HH:MM:SS,mmm, VTT HH:MM:SS.mmm, sequential SRT indices, WEBVTT header, hours>0, sub-second, multi-line - Edge cases: empty input, negative-time clamp, millisecond rounding, CRLF normalization - 21 unit tests in tests/test_subtitles.py; interrogate 100% Co-Authored-By: Claude Opus 4.8 Claude-Session: https://claude.ai/code/session_01CJRVbDrp1vGYkJgNHMGPpG --- subtitles.py | 155 ++++++++++++++++++++++++++++++++++ tests/test_subtitles.py | 179 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 334 insertions(+) create mode 100644 subtitles.py create mode 100644 tests/test_subtitles.py diff --git a/subtitles.py b/subtitles.py new file mode 100644 index 0000000..8042c9e --- /dev/null +++ b/subtitles.py @@ -0,0 +1,155 @@ +"""Subtitle and caption export for timestamped transcript segments. + +This module converts timestamped transcript segments into standard +subtitle formats: SubRip (``.srt``) and WebVTT (``.vtt``). It is fully +self-contained and depends only on the Python standard library, so it can +be used independently of the transcription pipeline. + +The input contract is intentionally small: any object exposing ``start`` +and ``end`` (floats, in seconds) and ``text`` (a string) is accepted. A +lightweight :class:`Segment` dataclass is provided for convenience, but +duck-typed objects work equally well. +""" + +from dataclasses import dataclass +from pathlib import Path +from typing import Iterable, List, Union + + +@dataclass +class Segment: + """A single timestamped transcript segment. + + Attributes: + start: Segment start time, in seconds from the beginning of the media. + end: Segment end time, in seconds from the beginning of the media. + text: The caption text for this segment. May contain newlines for + multi-line captions. + """ + + start: float + end: float + text: str + + +def _format_timestamp(seconds: float, separator: str) -> str: + """Format a time offset as ``HH:MM:SSmmm``. + + Args: + seconds: The time offset in seconds. Negative values are clamped to + zero so that malformed input never produces a negative timestamp. + separator: The separator placed between seconds and milliseconds. + Use ``","`` for SubRip and ``"."`` for WebVTT. + + Returns: + The formatted timestamp string, e.g. ``"01:02:03,004"``. + """ + if seconds < 0: + seconds = 0.0 + # Round to whole milliseconds to avoid floating-point drift. + total_milliseconds = int(round(seconds * 1000)) + hours, remainder = divmod(total_milliseconds, 3_600_000) + minutes, remainder = divmod(remainder, 60_000) + secs, milliseconds = divmod(remainder, 1000) + return f"{hours:02d}:{minutes:02d}:{secs:02d}{separator}{milliseconds:03d}" + + +def _normalize_text(text: str) -> str: + """Normalize caption text for cue output. + + Collapses Windows and old-Mac line endings to ``\\n`` and strips a single + trailing newline so cue blocks are separated cleanly by the caller. + + Args: + text: The raw caption text. + + Returns: + The normalized caption text. + """ + normalized = text.replace("\r\n", "\n").replace("\r", "\n") + return normalized.strip("\n") + + +def to_srt(segments: Iterable[Segment]) -> str: + """Render transcript segments as SubRip (``.srt``) text. + + Each cue is numbered sequentially starting at 1 and uses the SubRip + timestamp format ``HH:MM:SS,mmm``. Cues are separated by a blank line. + + Args: + segments: An iterable of objects exposing ``start`` and ``end`` + (floats, seconds) and ``text`` (a string). + + Returns: + The complete SubRip document as a string. An empty iterable yields an + empty string. + """ + blocks: List[str] = [] + for index, segment in enumerate(segments, start=1): + start = _format_timestamp(segment.start, ",") + end = _format_timestamp(segment.end, ",") + text = _normalize_text(segment.text) + blocks.append(f"{index}\n{start} --> {end}\n{text}") + if not blocks: + return "" + return "\n\n".join(blocks) + "\n" + + +def to_vtt(segments: Iterable[Segment]) -> str: + """Render transcript segments as WebVTT (``.vtt``) text. + + The document begins with the mandatory ``WEBVTT`` header. Each cue uses + the WebVTT timestamp format ``HH:MM:SS.mmm`` and cues are separated by a + blank line. + + Args: + segments: An iterable of objects exposing ``start`` and ``end`` + (floats, seconds) and ``text`` (a string). + + Returns: + The complete WebVTT document as a string. Even with no segments the + ``WEBVTT`` header is always emitted. + """ + blocks: List[str] = [] + for segment in segments: + start = _format_timestamp(segment.start, ".") + end = _format_timestamp(segment.end, ".") + text = _normalize_text(segment.text) + blocks.append(f"{start} --> {end}\n{text}") + if not blocks: + return "WEBVTT\n" + return "WEBVTT\n\n" + "\n\n".join(blocks) + "\n" + + +def write_srt(segments: Iterable[Segment], path: Union[str, Path]) -> Path: + """Write transcript segments to a SubRip (``.srt``) file. + + Args: + segments: An iterable of objects exposing ``start``, ``end`` and + ``text``. + path: Destination file path. Parent directories are not created + automatically. + + Returns: + The :class:`~pathlib.Path` that was written. + """ + destination = Path(path) + destination.write_text(to_srt(segments), encoding="utf-8") + return destination + + +def write_vtt(segments: Iterable[Segment], path: Union[str, Path]) -> Path: + """Write transcript segments to a WebVTT (``.vtt``) file. + + Args: + segments: An iterable of objects exposing ``start``, ``end`` and + ``text``. + path: Destination file path. Parent directories are not created + automatically. + + Returns: + The :class:`~pathlib.Path` that was written. + """ + destination = Path(path) + destination.write_text(to_vtt(segments), encoding="utf-8") + return destination diff --git a/tests/test_subtitles.py b/tests/test_subtitles.py new file mode 100644 index 0000000..c2a2783 --- /dev/null +++ b/tests/test_subtitles.py @@ -0,0 +1,179 @@ +"""Tests for the subtitle/caption export module.""" + +import tempfile +import unittest +from pathlib import Path + +from subtitles import ( + Segment, + to_srt, + to_vtt, + write_srt, + write_vtt, +) + + +class TestToSrt(unittest.TestCase): + """Tests for :func:`subtitles.to_srt`.""" + + def test_empty_input_returns_empty_string(self): + """An empty segment list produces an empty SubRip document.""" + self.assertEqual(to_srt([]), "") + + def test_single_segment_format(self): + """A single cue is numbered 1 and uses comma-millisecond timestamps.""" + result = to_srt([Segment(0.0, 1.5, "Hello")]) + self.assertEqual( + result, + "1\n00:00:00,000 --> 00:00:01,500\nHello\n", + ) + + def test_sequential_indices(self): + """Cues are numbered sequentially starting at 1.""" + segments = [ + Segment(0.0, 1.0, "one"), + Segment(1.0, 2.0, "two"), + Segment(2.0, 3.0, "three"), + ] + result = to_srt(segments) + lines = result.splitlines() + # Index lines appear at the start of each cue block. + self.assertEqual(lines[0], "1") + self.assertEqual(lines[4], "2") + self.assertEqual(lines[8], "3") + + def test_sub_second_timestamps(self): + """Fractional seconds are rendered as milliseconds.""" + result = to_srt([Segment(0.001, 0.25, "x")]) + self.assertIn("00:00:00,001 --> 00:00:00,250", result) + + def test_hour_plus_timestamp(self): + """Offsets beyond one hour populate the hours field.""" + # 3661.789s == 01:01:01.789 + result = to_srt([Segment(3661.789, 3662.0, "later")]) + self.assertIn("01:01:01,789 --> 01:01:02,000", result) + + def test_multi_line_text_preserved(self): + """Newlines within caption text are preserved in the cue body.""" + result = to_srt([Segment(0.0, 1.0, "line one\nline two")]) + self.assertIn("line one\nline two", result) + + def test_uses_comma_separator_not_period(self): + """SubRip must use a comma before the milliseconds field.""" + result = to_srt([Segment(0.0, 1.0, "x")]) + self.assertIn("00:00:00,000", result) + self.assertNotIn("00:00:00.000", result) + + def test_cues_separated_by_blank_line(self): + """Consecutive cues are separated by exactly one blank line.""" + result = to_srt([Segment(0.0, 1.0, "a"), Segment(1.0, 2.0, "b")]) + self.assertIn("\na\n\n2\n", result) + + +class TestToVtt(unittest.TestCase): + """Tests for :func:`subtitles.to_vtt`.""" + + def test_empty_input_still_has_header(self): + """Even with no cues the WEBVTT header is emitted.""" + self.assertEqual(to_vtt([]), "WEBVTT\n") + + def test_header_present(self): + """A populated document begins with the WEBVTT header.""" + result = to_vtt([Segment(0.0, 1.0, "hi")]) + self.assertTrue(result.startswith("WEBVTT\n\n")) + + def test_single_segment_format(self): + """WebVTT cues use period-millisecond timestamps and no index.""" + result = to_vtt([Segment(0.0, 1.5, "Hello")]) + self.assertEqual( + result, + "WEBVTT\n\n00:00:00.000 --> 00:00:01.500\nHello\n", + ) + + def test_uses_period_separator_not_comma(self): + """WebVTT must use a period before the milliseconds field.""" + result = to_vtt([Segment(0.0, 1.0, "x")]) + self.assertIn("00:00:00.000", result) + self.assertNotIn("00:00:00,000", result) + + def test_hour_plus_timestamp(self): + """Offsets beyond one hour populate the hours field.""" + result = to_vtt([Segment(3661.789, 3662.0, "later")]) + self.assertIn("01:01:01.789 --> 01:01:02.000", result) + + def test_no_sequential_index_lines(self): + """WebVTT cues are not prefixed with numeric indices.""" + result = to_vtt([Segment(0.0, 1.0, "a"), Segment(1.0, 2.0, "b")]) + # The line immediately after the blank header line is a timestamp. + lines = result.splitlines() + self.assertEqual(lines[0], "WEBVTT") + self.assertEqual(lines[1], "") + self.assertIn("-->", lines[2]) + + +class TestTimestampEdgeCases(unittest.TestCase): + """Edge-case tests for timestamp formatting.""" + + def test_negative_time_clamped_to_zero(self): + """Negative offsets are clamped rather than producing bad output.""" + result = to_srt([Segment(-5.0, 1.0, "x")]) + self.assertIn("00:00:00,000 --> 00:00:01,000", result) + + def test_millisecond_rounding(self): + """Values are rounded to the nearest whole millisecond.""" + # 1.0004s rounds to 000ms, 1.0006s rounds to 001ms. + self.assertIn("00:00:01,000", to_srt([Segment(1.0004, 2.0, "a")])) + self.assertIn("00:00:01,001", to_srt([Segment(1.0006, 2.0, "a")])) + + def test_crlf_normalized(self): + """Windows line endings in text are normalized to newlines.""" + result = to_srt([Segment(0.0, 1.0, "a\r\nb")]) + self.assertIn("a\nb", result) + self.assertNotIn("\r", result) + + def test_duck_typed_object_accepted(self): + """Any object with start/end/text attributes works.""" + + class Cue: + """Minimal duck-typed segment stand-in.""" + + start = 0.0 + end = 1.0 + text = "duck" + + result = to_srt([Cue()]) + self.assertIn("duck", result) + self.assertIn("00:00:00,000 --> 00:00:01,000", result) + + +class TestFileWriters(unittest.TestCase): + """Tests for the file-writing helpers.""" + + def test_write_srt_roundtrip(self): + """write_srt writes exactly what to_srt produces.""" + segments = [Segment(0.0, 1.0, "hi")] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.srt" + returned = write_srt(segments, path) + self.assertEqual(returned, path) + self.assertEqual(path.read_text(encoding="utf-8"), to_srt(segments)) + + def test_write_vtt_roundtrip(self): + """write_vtt writes exactly what to_vtt produces.""" + segments = [Segment(0.0, 1.0, "hi")] + with tempfile.TemporaryDirectory() as tmp: + path = Path(tmp) / "out.vtt" + returned = write_vtt(segments, path) + self.assertEqual(returned, path) + self.assertEqual(path.read_text(encoding="utf-8"), to_vtt(segments)) + + def test_write_srt_accepts_string_path(self): + """A string path is accepted and returned as a Path.""" + with tempfile.TemporaryDirectory() as tmp: + path_str = str(Path(tmp) / "out.srt") + returned = write_srt([Segment(0.0, 1.0, "x")], path_str) + self.assertTrue(returned.exists()) + + +if __name__ == "__main__": + unittest.main()