182 lines
7.5 KiB
Python
182 lines
7.5 KiB
Python
# Copyright 2026 Umbra Tools (umbratools.com)
|
||
# SPDX-License-Identifier: Apache-2.0
|
||
|
||
from __future__ import annotations
|
||
|
||
from umbraaudio.mapping.segmenter import Segmenter
|
||
|
||
|
||
def test_multisentence_quote_stays_one_utterance() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"“Fifty light-years. An old freighter like this could only jump "
|
||
"twenty-five light-years. How can Deal Hunter do that?”"
|
||
)
|
||
utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert len(dialogues) == 1
|
||
assert dialogues[0]["text"].endswith("?”")
|
||
assert "How can Deal Hunter do that?" in dialogues[0]["text"]
|
||
|
||
|
||
def test_quote_with_inner_single_quotes_not_mangled() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"“Our best analysis would indicate that it is a message. ‘Don’t help "
|
||
"Princess Kainda.’ Princess, we must consider how to deal with you now.”"
|
||
)
|
||
utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert len(dialogues) == 1
|
||
assert "‘Don’t help Princess Kainda.’" in dialogues[0]["text"]
|
||
assert "Princess, we must consider" in dialogues[0]["text"]
|
||
|
||
|
||
def test_dialogue_and_narration_split_kept_separate() -> None:
|
||
seg = Segmenter()
|
||
text = "Nanny removed her helmet. “How are you feeling, Princess?”"
|
||
utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
|
||
types = [u["type"] for u in utts]
|
||
assert "narration" in types
|
||
assert "dialogue" in types
|
||
|
||
|
||
def test_carryover_speaker_after_tag() -> None:
|
||
seg = Segmenter()
|
||
text = "Nanny removed her helmet. “How are you feeling, Princess?”\n“I have a headache.”"
|
||
utts = seg.segment(text, 1, character_names=["Nanny", "Kainda", "Princess"])
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert dialogues[0]["character_name"] == "Nanny"
|
||
assert dialogues[1]["character_name"] == "Princess"
|
||
|
||
|
||
def test_outer_quote_closed_when_inner_quote_at_end() -> None:
|
||
seg = Segmenter()
|
||
text = "“I am Support Robot 74532, but I have been named ‘Nanny.’”"
|
||
utts = seg.segment(text, 1, character_names=["Nanny", "Kainda"])
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert len(dialogues) == 1
|
||
assert dialogues[0]["text"].endswith("’”")
|
||
assert "‘Nanny.’" in dialogues[0]["text"]
|
||
|
||
|
||
def test_inner_quote_kept_when_quote_continues() -> None:
|
||
seg = Segmenter()
|
||
text = "“She asked me to say, ‘Nothing will happen to you on my ship.’”"
|
||
utts = seg.segment(text, 1, character_names=["Kainda"])
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert len(dialogues) == 1
|
||
assert dialogues[0]["text"] == text
|
||
|
||
|
||
def test_author_note_splits_into_sentences() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"Author’s Note\n\n"
|
||
"This novel, Deal Hunter, does not use the terms android or droid. "
|
||
"Google uses “Android” to refer to the operating system on phones.\n\n"
|
||
"Deal Hunter uses the generic terms of robot and artificial intelligence (AI)."
|
||
)
|
||
utts = seg.segment(text, 5, chapter_title="Author's Note")
|
||
texts = [u["text"] for u in utts]
|
||
assert len(texts) == 4
|
||
assert texts[0] == "Author’s Note"
|
||
assert "does not use the terms android or droid." in texts[1]
|
||
assert "Google uses “Android” to refer to the operating system on phones." in texts[2]
|
||
assert texts[3].startswith("Deal Hunter uses the generic terms")
|
||
|
||
|
||
def test_author_note_numbered_list_splits_on_sentence_end() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"Author’s Note\n\n"
|
||
"This novel started as an exploration of several questions.\n\n"
|
||
"1) A dragon buys nothing, yet it hoards gold. Why does the dragon"
|
||
" consider gold plates, cups, and ornaments a treasure?\n"
|
||
"2) How does a dwarven city inside the mountains feed everyone with no farms?\n"
|
||
"3) Why is the city located inside that specific mountain?\n"
|
||
"4) Dragons and dwarves live in the mountains. Hmm. Is there more?"
|
||
)
|
||
utts = seg.segment(text, 5, chapter_title="Author's Note")
|
||
texts = [u["text"] for u in utts]
|
||
question = "Why does the dragon consider gold plates, cups, and ornaments a treasure?"
|
||
item2 = "2) How does a dwarven city inside the mountains feed everyone with no farms?"
|
||
item3 = "3) Why is the city located inside that specific mountain?"
|
||
assert any(t == question for t in texts)
|
||
assert any(t == item2 for t in texts)
|
||
assert any(t == item3 for t in texts)
|
||
assert any(t == "4) Dragons and dwarves live in the mountains." for t in texts)
|
||
assert any(t == "Is there more?" for t in texts)
|
||
assert all(u["type"] == "narration" for u in utts)
|
||
|
||
|
||
def test_sentence_end_does_not_split_on_money_or_number() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"The price is $2. 50 more than the last model. Section 3. 4 of the manual"
|
||
" explains it all."
|
||
)
|
||
utts = seg.segment(text, 1, character_names=["Nanny"])
|
||
narration = [u for u in utts if u["type"] == "narration"]
|
||
joined = " ".join(u["text"] for u in narration)
|
||
assert "$2. 50 more than the last model" in joined
|
||
assert "Section 3. 4 of the manual" in joined
|
||
|
||
|
||
def test_author_note_all_narration() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"Author’s Note\n\n"
|
||
"Every star has humans, and they are looking for resources. "
|
||
"The Dial sector has several pirate clans."
|
||
)
|
||
utts = seg.segment(text, 5, chapter_title="Author's Note")
|
||
assert all(u["type"] == "narration" for u in utts)
|
||
assert all(u["character_name"] == "narrator" for u in utts)
|
||
|
||
|
||
def test_author_note_merges_abbreviations() -> None:
|
||
seg = Segmenter()
|
||
text = "Author’s Note\n\nThis is a work of fiction. Dr. Smith wrote it."
|
||
utts = seg.segment(text, 5, chapter_title="Author's Note")
|
||
texts = [u["text"] for u in utts]
|
||
assert "This is a work of fiction." in texts
|
||
assert "Dr. Smith wrote it." in texts
|
||
|
||
|
||
def test_unclosed_quote_does_not_swallow_following_narration() -> None:
|
||
seg = Segmenter()
|
||
text = (
|
||
"As she says this, there is a roar, and Niñera says, “The captain is "
|
||
"changing course, and that was a missile he is firing at the pirates. \n"
|
||
"As Niñera helps Kainda, they feel the impacts more, and the air "
|
||
"pressure alarm sounds. The computer voice announces, “Hull breach. "
|
||
"Please put on emergency space suits.”\nSealing the suit, the robot "
|
||
"picks up the helmet."
|
||
)
|
||
utts = seg.segment(text, 6, character_names=["Niñera", "Kainda"])
|
||
assert len(utts) >= 3
|
||
joined = " ".join(u["text"] for u in utts)
|
||
assert "Hull breach" in joined
|
||
hull = [u for u in utts if "Hull breach" in u["text"]]
|
||
assert len(hull) == 1
|
||
assert "changing course" not in hull[0]["text"]
|
||
|
||
|
||
def test_unterminated_quote_at_eof_becomes_plain_text() -> None:
|
||
seg = Segmenter()
|
||
text = "She stared out the viewport. “We are never going to make it"
|
||
utts = seg.segment(text, 1)
|
||
assert all(u["type"] == "narration" for u in utts)
|
||
assert any("never going to make it" in u["text"] for u in utts)
|
||
|
||
|
||
def test_balanced_multisentence_quote_still_intact() -> None:
|
||
seg = Segmenter()
|
||
text = "He said, “Line one. Line two. Line three.” Then he left."
|
||
utts = seg.segment(text, 1)
|
||
dialogues = [u for u in utts if u["type"] == "dialogue"]
|
||
assert len(dialogues) == 1
|
||
assert dialogues[0]["text"].startswith("“")
|
||
assert dialogues[0]["text"].endswith("”")
|