diff --git a/abogen/domain/subtitle_generation.py b/abogen/domain/subtitle_generation.py new file mode 100644 index 0000000..99500dc --- /dev/null +++ b/abogen/domain/subtitle_generation.py @@ -0,0 +1,358 @@ +"""Subtitle generation utilities for audiobook generation. + +This module provides functions for processing TTS tokens into subtitle entries +according to various subtitle modes (Line, Sentence, Sentence + Comma, +Sentence + Highlighting). +""" + +from __future__ import annotations + +import re +from typing import List, Optional, Tuple + + +# Punctuation constants for sentence splitting +PUNCTUATION_SENTENCE = ".!?\u061f\u3002\uff01\uff1f" # .!? .?. ?? +PUNCTUATION_SENTENCE_COMMA = ".!?,\u3001\u061f\u3002\uff01\uff0c\uff1f" # .!?, ,. ?? + + +def process_subtitle_tokens( + tokens_with_timestamps: List[dict], + subtitle_entries: List[Tuple[float, float, str]], + max_subtitle_words: int, + subtitle_mode: str, + lang_code: str, + use_spacy_segmentation: bool = False, + fallback_end_time: Optional[float] = None, +) -> None: + """Process TTS tokens into subtitle entries according to the subtitle mode. + + This function modifies subtitle_entries in-place by appending new entries. + + Args: + tokens_with_timestamps: List of token dictionaries with 'start', 'end', 'text', + and 'whitespace' keys. + subtitle_entries: List to append subtitle entries to (modified in-place). + Each entry is a tuple of (start_time, end_time, text). + max_subtitle_words: Maximum number of words per subtitle entry. + subtitle_mode: One of "Disabled", "Line", "Sentence", "Sentence + Comma", + "Sentence + Highlighting", or a string like "5" for word-count mode. + lang_code: Language code for spaCy processing (e.g., "a" for English). + use_spacy_segmentation: Whether to use spaCy for sentence boundary detection. + fallback_end_time: Fallback end time for the last entry if none is available. + """ + if not tokens_with_timestamps: + return + + processed_tokens = tokens_with_timestamps + + # For English with spaCy enabled and sentence-based modes, use spaCy for sentence boundaries + # spaCy is disabled when subtitle mode is "Disabled" or "Line" + use_spacy_for_english = ( + use_spacy_segmentation + and subtitle_mode not in ["Disabled", "Line"] + and lang_code in ["a", "b"] + and subtitle_mode in ["Sentence", "Sentence + Comma"] + ) + + if subtitle_mode == "Sentence + Highlighting": + _process_karaoke_highlighting( + processed_tokens, subtitle_entries, max_subtitle_words, fallback_end_time + ) + elif subtitle_mode in ["Sentence", "Sentence + Comma", "Line"]: + if use_spacy_for_english and subtitle_mode != "Line": + _process_spacy_sentences( + processed_tokens, subtitle_entries, max_subtitle_words, + subtitle_mode, lang_code, fallback_end_time + ) + else: + _process_regex_sentences( + processed_tokens, subtitle_entries, max_subtitle_words, + subtitle_mode, fallback_end_time + ) + else: + # Word count-based grouping (e.g., "5" for 5-word groups) + _process_word_count( + processed_tokens, subtitle_entries, max_subtitle_words, + subtitle_mode, fallback_end_time + ) + + +def _process_karaoke_highlighting( + tokens: List[dict], + subtitle_entries: List[Tuple[float, float, str]], + max_subtitle_words: int, + fallback_end_time: Optional[float], +) -> None: + """Process tokens for Sentence + Highlighting mode (karaoke effect).""" + separator = rf"[{re.escape(PUNCTUATION_SENTENCE)}]" + current_sentence = [] + word_count = 0 + + for token in tokens: + current_sentence.append(token) + word_count += 1 + + # Split sentences based on separator or word count + if ( + re.search(separator, token["text"]) and token.get("whitespace") == " " + ) or word_count >= max_subtitle_words: + if current_sentence: + # Create karaoke subtitle entry for this sentence + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + + # Generate karaoke text with timing + karaoke_text = "" + for t in current_sentence: + # Calculate duration in centiseconds + duration = ( + t["end"] - t["start"] + if t.get("end") is not None and t.get("start") is not None + else 0.5 + ) + duration_cs = int(duration * 100) + # Add karaoke effect + karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}" + + subtitle_entries.append( + (start_time, end_time, karaoke_text.strip()) + ) + current_sentence = [] + word_count = 0 + + # Add any remaining tokens as a sentence + if current_sentence: + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + + # Generate karaoke text for remaining tokens + karaoke_text = "" + for t in current_sentence: + duration = t["end"] - t["start"] if t.get("end") and t.get("start") else 0.5 + duration_cs = int(duration * 100) + karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}" + subtitle_entries.append((start_time, end_time, karaoke_text.strip())) + + # Fallback for last entry + _apply_fallback_end_time(subtitle_entries, fallback_end_time) + + +def _process_spacy_sentences( + tokens: List[dict], + subtitle_entries: List[Tuple[float, float, str]], + max_subtitle_words: int, + subtitle_mode: str, + lang_code: str, + fallback_end_time: Optional[float], +) -> None: + """Process tokens using spaCy for sentence boundary detection.""" + try: + from abogen.spacy_utils import get_spacy_model + except ImportError: + # Fall back to regex if spaCy is not available + _process_regex_sentences( + tokens, subtitle_entries, max_subtitle_words, + subtitle_mode, fallback_end_time + ) + return + + nlp = get_spacy_model(lang_code) + if not nlp: + _process_regex_sentences( + tokens, subtitle_entries, max_subtitle_words, + subtitle_mode, fallback_end_time + ) + return + + # Build full text and track character positions to token indices + full_text = "" + for token in tokens: + text_part = token["text"] + (token.get("whitespace") or "") + full_text += text_part + + # Get sentence boundaries from spaCy + doc = nlp(full_text) + sentence_boundaries = [sent.end_char for sent in doc.sents] + + # For "Sentence + Comma" mode, also split on commas + if subtitle_mode == "Sentence + Comma": + comma_positions = [ + i + 1 for i, c in enumerate(full_text) if c == "," + ] + sentence_boundaries = sorted( + set(sentence_boundaries + comma_positions) + ) + + # Group tokens by sentence boundaries + current_sentence = [] + word_count = 0 + current_char_pos = 0 + boundary_idx = 0 + + for token in tokens: + current_sentence.append(token) + word_count += 1 + text_len = len(token["text"]) + len(token.get("whitespace") or "") + current_char_pos += text_len + + # Check if we've hit a sentence boundary or max words + at_boundary = ( + boundary_idx < len(sentence_boundaries) + and current_char_pos >= sentence_boundaries[boundary_idx] + ) + if at_boundary or word_count >= max_subtitle_words: + if current_sentence: + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + sentence_text = "".join( + t["text"] + (t.get("whitespace") or "") + for t in current_sentence + ) + subtitle_entries.append( + (start_time, end_time, sentence_text.strip()) + ) + current_sentence = [] + word_count = 0 + if at_boundary: + boundary_idx += 1 + + # Add remaining tokens + if current_sentence: + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + sentence_text = "".join( + t["text"] + (t.get("whitespace") or "") + for t in current_sentence + ) + subtitle_entries.append( + (start_time, end_time, sentence_text.strip()) + ) + + # Fallback for last entry + _apply_fallback_end_time(subtitle_entries, fallback_end_time) + + +def _process_regex_sentences( + tokens: List[dict], + subtitle_entries: List[Tuple[float, float, str]], + max_subtitle_words: int, + subtitle_mode: str, + fallback_end_time: Optional[float], +) -> None: + """Process tokens using regex for sentence boundary detection.""" + # Define separator pattern based on mode + if subtitle_mode == "Line": + separator = r"\n" + elif subtitle_mode == "Sentence": + # Use punctuation without comma + separator = rf"[{re.escape(PUNCTUATION_SENTENCE)}]" + else: # Sentence + Comma + # Use punctuation with comma + separator = rf"[{re.escape(PUNCTUATION_SENTENCE_COMMA)}]" + + current_sentence = [] + word_count = 0 + + for token in tokens: + current_sentence.append(token) + word_count += 1 + + # Split sentences based on separator or word count + if ( + re.search(separator, token["text"]) and token.get("whitespace") == " " + ) or word_count >= max_subtitle_words: + if current_sentence: + # Create subtitle entry for this sentence + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + + # Simplified text joining logic + sentence_text = "" + for t in current_sentence: + sentence_text += t["text"] + (t.get("whitespace") or "") + + subtitle_entries.append( + (start_time, end_time, sentence_text.strip()) + ) + current_sentence = [] + word_count = 0 + + # Add any remaining tokens as a sentence + if current_sentence: + start_time = current_sentence[0]["start"] + end_time = current_sentence[-1]["end"] + + # Simplified text joining logic + sentence_text = "" + for t in current_sentence: + sentence_text += t["text"] + (t.get("whitespace") or "") + subtitle_entries.append((start_time, end_time, sentence_text.strip())) + + # Fallback for last entry + _apply_fallback_end_time(subtitle_entries, fallback_end_time) + + +def _process_word_count( + tokens: List[dict], + subtitle_entries: List[Tuple[float, float, str]], + max_subtitle_words: int, + subtitle_mode: str, + fallback_end_time: Optional[float], +) -> None: + """Process tokens by counting spaces (word count mode).""" + try: + word_count = int(subtitle_mode.split()[0]) + word_count = min(word_count, max_subtitle_words) + except (ValueError, IndexError): + word_count = 1 + + current_group = [] + space_count = 0 + + for token in tokens: + current_group.append(token) + + # Count spaces after tokens (in the whitespace field) + if token.get("whitespace", "") == " ": + space_count += 1 + + # Split after counting N spaces + if space_count >= word_count: + text = "".join( + t["text"] + (t.get("whitespace") or "") + for t in current_group + ) + subtitle_entries.append( + ( + current_group[0]["start"], + current_group[-1]["end"], + text.strip(), + ) + ) + current_group = [] + space_count = 0 + + # Add any remaining tokens + if current_group: + text = "".join( + t["text"] + (t.get("whitespace") or "") for t in current_group + ) + subtitle_entries.append( + (current_group[0]["start"], current_group[-1]["end"], text.strip()) + ) + + # Fallback for last entry + _apply_fallback_end_time(subtitle_entries, fallback_end_time) + + +def _apply_fallback_end_time( + subtitle_entries: List[Tuple[float, float, str]], + fallback_end_time: Optional[float], +) -> None: + """Apply fallback end time to the last entry if needed.""" + if subtitle_entries and fallback_end_time is not None: + last_entry = subtitle_entries[-1] + start, end, text = last_entry + if end is None or end <= start or end <= 0: + subtitle_entries[-1] = (start, fallback_end_time, text) diff --git a/abogen/pyqt/conversion.py b/abogen/pyqt/conversion.py index 6ea017d..fcfdbbc 100644 --- a/abogen/pyqt/conversion.py +++ b/abogen/pyqt/conversion.py @@ -35,6 +35,7 @@ from abogen.domain.audio_buffer import ( normalize_audio, SAMPLE_RATE, ) +from abogen.domain.subtitle_generation import process_subtitle_tokens import abogen.hf_tracker as hf_tracker import static_ffmpeg import threading # for efficient waiting @@ -1997,270 +1998,16 @@ class ConversionThread(QThread): fallback_end_time=None, ): """Helper function to process subtitle tokens according to the subtitle mode""" - if not tokens_with_timestamps: - return - - processed_tokens = tokens_with_timestamps # Use tokens directly - - # For English with spaCy enabled and sentence-based modes, use spaCy for sentence boundaries - # spaCy is disabled when subtitle mode is "Disabled" or "Line" - use_spacy_for_english = ( - getattr(self, "use_spacy_segmentation", False) - and self.subtitle_mode not in ["Disabled", "Line"] - and self.lang_code in ["a", "b"] - and self.subtitle_mode in ["Sentence", "Sentence + Comma"] + process_subtitle_tokens( + tokens_with_timestamps=tokens_with_timestamps, + subtitle_entries=subtitle_entries, + max_subtitle_words=max_subtitle_words, + subtitle_mode=self.subtitle_mode, + lang_code=self.lang_code, + use_spacy_segmentation=getattr(self, "use_spacy_segmentation", False), + fallback_end_time=fallback_end_time, ) - # Use processed_tokens instead of tokens_with_timestamps for the rest of the method - if self.subtitle_mode == "Sentence + Highlighting": - # Sentence-based processing with karaoke highlighting - # Use punctuation without comma - separator = r"[{}]".format(self.PUNCTUATION_SENTENCE) - current_sentence = [] - word_count = 0 - - for token in processed_tokens: # Updated to use processed_tokens - current_sentence.append(token) - word_count += 1 - - # Split sentences based on separator or word count - if ( - re.search(separator, token["text"]) and token["whitespace"] == " " - ) or word_count >= max_subtitle_words: - if current_sentence: - # Create karaoke subtitle entry for this sentence - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - - # Generate karaoke text with background highlighting - karaoke_text = "" - for t in current_sentence: - # Calculate duration in centiseconds - duration = ( - t["end"] - t["start"] - if t["end"] and t["start"] - else 0.5 - ) - duration_cs = int(duration * 100) - # Add karaoke effect - relies on style's SecondaryColour for highlighting - karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}" - - subtitle_entries.append( - (start_time, end_time, karaoke_text.strip()) - ) - current_sentence = [] - word_count = 0 - - # Add any remaining tokens as a sentence - if current_sentence: - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - - # Generate karaoke text for remaining tokens - karaoke_text = "" - for t in current_sentence: - duration = t["end"] - t["start"] if t["end"] and t["start"] else 0.5 - duration_cs = int(duration * 100) - karaoke_text += f"{{\\kf{duration_cs}}}{t['text']}{t.get('whitespace', '') or ''}" - subtitle_entries.append((start_time, end_time, karaoke_text.strip())) - - # Fallback for last entry - if subtitle_entries and fallback_end_time is not None: - last_entry = subtitle_entries[-1] - start, end, text = last_entry - if end is None or end <= start or end <= 0: - subtitle_entries[-1] = (start, fallback_end_time, text) - - elif self.subtitle_mode in ["Sentence", "Sentence + Comma", "Line"]: - # Check if we should use spaCy for English sentence boundaries - if use_spacy_for_english and self.subtitle_mode != "Line": - # Use spaCy for English sentence boundary detection (model already loaded) - from abogen.spacy_utils import get_spacy_model - - nlp = get_spacy_model( - self.lang_code - ) # No log_callback since model is already loaded - if nlp: - # Build full text and track character positions to token indices - full_text = "" - char_to_token = [] # Maps character index to token index - for idx, token in enumerate(processed_tokens): - start_char = len(full_text) - text_part = token["text"] + (token.get("whitespace", "") or "") - full_text += text_part - char_to_token.extend([idx] * len(text_part)) - - # Get sentence boundaries from spaCy - doc = nlp(full_text) - sentence_boundaries = [sent.end_char for sent in doc.sents] - - # For "Sentence + Comma" mode, also split on commas - if self.subtitle_mode == "Sentence + Comma": - comma_positions = [ - i + 1 for i, c in enumerate(full_text) if c == "," - ] - sentence_boundaries = sorted( - set(sentence_boundaries + comma_positions) - ) - - # Group tokens by sentence boundaries - current_sentence = [] - word_count = 0 - current_char_pos = 0 - boundary_idx = 0 - - for idx, token in enumerate(processed_tokens): - current_sentence.append(token) - word_count += 1 - text_len = len(token["text"]) + len( - token.get("whitespace", "") or "" - ) - current_char_pos += text_len - - # Check if we've hit a sentence boundary or max words - at_boundary = ( - boundary_idx < len(sentence_boundaries) - and current_char_pos >= sentence_boundaries[boundary_idx] - ) - if at_boundary or word_count >= max_subtitle_words: - if current_sentence: - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - sentence_text = "".join( - t["text"] + (t.get("whitespace", "") or "") - for t in current_sentence - ) - subtitle_entries.append( - (start_time, end_time, sentence_text.strip()) - ) - current_sentence = [] - word_count = 0 - if at_boundary: - boundary_idx += 1 - - # Add remaining tokens - if current_sentence: - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - sentence_text = "".join( - t["text"] + (t.get("whitespace", "") or "") - for t in current_sentence - ) - subtitle_entries.append( - (start_time, end_time, sentence_text.strip()) - ) - - # Fallback for last entry - if subtitle_entries and fallback_end_time is not None: - last_entry = subtitle_entries[-1] - start, end, text = last_entry - if end is None or end <= start or end <= 0: - subtitle_entries[-1] = (start, fallback_end_time, text) - return # Exit early, spaCy processing complete - - # Default regex-based processing (non-English or spaCy unavailable) - # Define separator pattern based on mode - if self.subtitle_mode == "Line": - separator = r"\n" - elif self.subtitle_mode == "Sentence": - # Use punctuation without comma - separator = r"[{}]".format(self.PUNCTUATION_SENTENCE) - else: # Sentence + Comma - # Use punctuation with comma - separator = r"[{}]".format(self.PUNCTUATION_SENTENCE_COMMA) - current_sentence = [] - word_count = 0 - - for token in processed_tokens: # Updated to use processed_tokens - current_sentence.append(token) - word_count += 1 - - # Split sentences based on separator or word count - if ( - re.search(separator, token["text"]) and token["whitespace"] == " " - ) or word_count >= max_subtitle_words: - if current_sentence: - # Create subtitle entry for this sentence - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - - # Simplified text joining logic - sentence_text = "" - for t in current_sentence: - sentence_text += t["text"] + (t.get("whitespace", "") or "") - - subtitle_entries.append( - (start_time, end_time, sentence_text.strip()) - ) - current_sentence = [] - word_count = 0 - - # Add any remaining tokens as a sentence - if current_sentence: - start_time = current_sentence[0]["start"] - end_time = current_sentence[-1]["end"] - - # Simplified text joining logic - sentence_text = "" - for t in current_sentence: - sentence_text += t["text"] + (t.get("whitespace", "") or "") - subtitle_entries.append((start_time, end_time, sentence_text.strip())) - - # Fallback for last entry - if subtitle_entries and fallback_end_time is not None: - last_entry = subtitle_entries[-1] - start, end, text = last_entry - if end is None or end <= start or end <= 0: - subtitle_entries[-1] = (start, fallback_end_time, text) - - else: - # Word count-based grouping - simply count spaces and split after N spaces - try: - word_count = int(self.subtitle_mode.split()[0]) - word_count = min(word_count, max_subtitle_words) - except (ValueError, IndexError): - word_count = 1 - - current_group = [] - space_count = 0 - - for token in processed_tokens: - current_group.append(token) - - # Count spaces after tokens (in the whitespace field) - if token.get("whitespace", "") == " ": - space_count += 1 - - # Split after counting N spaces - if space_count >= word_count: - text = "".join( - t["text"] + (t.get("whitespace", "") or "") - for t in current_group - ) - subtitle_entries.append( - ( - current_group[0]["start"], - current_group[-1]["end"], - text.strip(), - ) - ) - current_group = [] - space_count = 0 - - # Add any remaining tokens - if current_group: - text = "".join( - t["text"] + (t.get("whitespace", "") or "") for t in current_group - ) - subtitle_entries.append( - (current_group[0]["start"], current_group[-1]["end"], text.strip()) - ) - - # Fallback for last entry - if subtitle_entries and fallback_end_time is not None: - last_entry = subtitle_entries[-1] - start, end, text = last_entry - if end is None or end <= start or end <= 0: - subtitle_entries[-1] = (start, fallback_end_time, text) + return def cancel(self): self.cancel_requested = True diff --git a/tests/test_subtitle_generation.py b/tests/test_subtitle_generation.py new file mode 100644 index 0000000..715f4d8 --- /dev/null +++ b/tests/test_subtitle_generation.py @@ -0,0 +1,195 @@ +"""Tests for abogen.domain.subtitle_generation module.""" + +import pytest + +from abogen.domain.subtitle_generation import ( + process_subtitle_tokens, + PUNCTUATION_SENTENCE, + PUNCTUATION_SENTENCE_COMMA, +) + + +class TestProcessSubtitleTokens: + """Tests for process_subtitle_tokens function.""" + + def test_empty_tokens(self): + """Test processing empty token list does nothing.""" + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=[], + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Sentence", + lang_code="a", + ) + assert entries == [] + + def test_disabled_mode(self): + """Test Disabled mode returns no entries.""" + tokens = [ + {"start": 0.0, "end": 1.0, "text": "Hello", "whitespace": " "}, + {"start": 1.0, "end": 2.0, "text": "world", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Disabled", + lang_code="a", + ) + assert entries == [] + + def test_line_mode_basic(self): + """Test Line mode splits on newlines.""" + tokens = [ + {"start": 0.0, "end": 1.0, "text": "First line", "whitespace": "\n"}, + {"start": 1.0, "end": 2.0, "text": "Second line", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Line", + lang_code="a", + ) + assert len(entries) == 2 + assert entries[0][2] == "First line" + assert entries[1][2] == "Second line" + + def test_sentence_mode_punctuation_split(self): + """Test Sentence mode splits on sentence punctuation.""" + tokens = [ + {"start": 0.0, "end": 0.5, "text": "First sentence", "whitespace": " "}, + {"start": 0.5, "end": 1.0, "text": ".", "whitespace": " "}, + {"start": 1.0, "end": 1.5, "text": "Second sentence", "whitespace": " "}, + {"start": 1.5, "end": 2.0, "text": ".", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Sentence", + lang_code="a", + ) + assert len(entries) >= 1 + # Should have at least one entry with both sentences or split + combined_text = " ".join(e[2] for e in entries) + assert "First sentence" in combined_text + assert "Second sentence" in combined_text + + def test_word_count_mode(self): + """Test word count mode (e.g., '5' for 5 words per entry).""" + tokens = [ + {"start": 0.0, "end": 0.2, "text": "word1", "whitespace": " "}, + {"start": 0.2, "end": 0.4, "text": "word2", "whitespace": " "}, + {"start": 0.4, "end": 0.6, "text": "word3", "whitespace": " "}, + {"start": 0.6, "end": 0.8, "text": "word4", "whitespace": " "}, + {"start": 0.8, "end": 1.0, "text": "word5", "whitespace": " "}, + {"start": 1.0, "end": 1.2, "text": "word6", "whitespace": " "}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="2", # 2 words per entry + lang_code="a", + ) + assert len(entries) >= 2 + # Check that entries are split roughly by word count + for entry in entries: + # Each entry should have at least one word + assert len(entry[2].split()) >= 1 + + def test_fallback_end_time(self): + """Test fallback_end_time is applied when end time is invalid.""" + tokens = [ + {"start": 0.0, "end": None, "text": "Test", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Line", + lang_code="a", + fallback_end_time=10.0, + ) + assert len(entries) == 1 + assert entries[0][1] == 10.0 # Should use fallback + + def test_karaoke_highlighting_mode(self): + """Test Sentence + Highlighting mode generates karaoke tags.""" + tokens = [ + {"start": 0.0, "end": 0.5, "text": "Hello", "whitespace": " "}, + {"start": 0.5, "end": 1.0, "text": "world", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Sentence + Highlighting", + lang_code="a", + ) + assert len(entries) >= 1 + # Should contain karaoke tags + text = entries[0][2] + assert "{\\kf" in text + + def test_max_subtitle_words_limit(self): + """Test that max_subtitle_words limits entry length.""" + tokens = [ + {"start": float(i), "end": float(i + 0.1), "text": f"word{i}", "whitespace": " "} + for i in range(10) + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=3, + subtitle_mode="Line", + lang_code="a", + ) + # Should have more than 1 entry due to word limit + assert len(entries) > 1 + + def test_preserves_token_timing(self): + """Test that token timing is preserved in entries.""" + tokens = [ + {"start": 0.0, "end": 1.0, "text": "First", "whitespace": " "}, + {"start": 1.0, "end": 2.0, "text": "Second", "whitespace": ""}, + ] + entries = [] + process_subtitle_tokens( + tokens_with_timestamps=tokens, + subtitle_entries=entries, + max_subtitle_words=50, + subtitle_mode="Sentence", + lang_code="a", + ) + assert len(entries) >= 1 + # Check that timing is preserved + for entry in entries: + assert entry[0] >= 0.0 + assert entry[1] >= entry[0] + + +class TestPunctuationConstants: + """Tests for punctuation constants.""" + + def test_punctuation_sentence_contains_basic(self): + """Test PUNCTUATION_SENTENCE contains basic sentence punctuation.""" + assert "." in PUNCTUATION_SENTENCE + assert "!" in PUNCTUATION_SENTENCE + assert "?" in PUNCTUATION_SENTENCE + + def test_punctuation_sentence_comma_contains_comma(self): + """Test PUNCTUATION_SENTENCE_COMMA contains comma.""" + assert "," in PUNCTUATION_SENTENCE_COMMA + assert "." in PUNCTUATION_SENTENCE_COMMA + assert "!" in PUNCTUATION_SENTENCE_COMMA + assert "?" in PUNCTUATION_SENTENCE_COMMA