mirror of
https://github.com/denizsafak/abogen.git
synced 2026-07-18 13:40:27 +02:00
refactor: extract subtitle token processing to domain layer
- Add abogen/domain/subtitle_generation.py with: - process_subtitle_tokens(): main function for converting TTS tokens to subtitles - Support for all subtitle modes: Line, Sentence, Sentence + Comma, Sentence + Highlighting - Support for word-count based grouping (e.g., '5' for 5 words per entry) - spaCy integration for English sentence boundary detection - Karaoke highlighting tags for Sentence + Highlighting mode - Punctuation constants for sentence splitting - Update abogen/pyqt/conversion.py: - Replace _process_subtitle_tokens method body with call to domain function - Remove ~260 lines of duplicate logic - Add tests/test_subtitle_generation.py with comprehensive unit tests
This commit is contained in:
@@ -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
|
||||
Reference in New Issue
Block a user