diff --git a/tests/test_word_solve.py b/tests/test_word_solve.py new file mode 100644 index 0000000..f383e51 --- /dev/null +++ b/tests/test_word_solve.py @@ -0,0 +1,88 @@ +# tests/test_word_solve.py +"""Tests for word dictionary solve mode.""" + +import subprocess +import types +import pytest +from PIL import ImageFont + +from unredact.pipeline.dictionary import solve_word_dictionary + + +def _get_test_font() -> ImageFont.FreeTypeFont: + result = subprocess.run( + ["fc-match", "--format=%{file}", "Liberation Serif"], + capture_output=True, text=True, + ) + path = result.stdout.strip() + return ImageFont.truetype(path, 40) + + +class TestSolveWordDictionarySingleWord: + def test_finds_exact_noun(self): + font = _get_test_font() + target = font.getlength("house") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="lowercase", + )) + texts = [r.text for r in results] + assert "house" in texts + + def test_plural_mode_filters(self): + font = _get_test_font() + target = font.getlength("houses") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="lowercase", ensure_plural=True, + )) + texts = [r.text for r in results] + assert "houses" in texts + + def test_casing_capitalized(self): + font = _get_test_font() + target = font.getlength("House") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="capitalized", + )) + texts = [r.text for r in results] + assert "House" in texts + + def test_casing_uppercase(self): + font = _get_test_font() + target = font.getlength("HOUSE") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="uppercase", + )) + texts = [r.text for r in results] + assert "HOUSE" in texts + + def test_known_start_filters(self): + font = _get_test_font() + target = font.getlength("house") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="lowercase", known_start="ho", + )) + texts = [r.text for r in results] + assert "house" in texts + assert all(t.lower().startswith("ho") for t in texts) + + def test_known_end_filters(self): + font = _get_test_font() + target = font.getlength("house") + results = list(solve_word_dictionary( + font, target, tolerance=0.5, + casing="lowercase", known_end="se", + )) + texts = [r.text for r in results] + assert "house" in texts + assert all(t.lower().endswith("se") for t in texts) + + def test_returns_generator(self): + font = _get_test_font() + target = font.getlength("dog") + gen = solve_word_dictionary(font, target, tolerance=0.5, casing="lowercase") + assert isinstance(gen, types.GeneratorType) diff --git a/unredact/pipeline/dictionary.py b/unredact/pipeline/dictionary.py index 6b4b6d2..fe75878 100644 --- a/unredact/pipeline/dictionary.py +++ b/unredact/pipeline/dictionary.py @@ -263,3 +263,57 @@ def solve_name_dictionary( results.sort(key=lambda r: (r.error, r.text)) return results + + +def solve_word_dictionary( + font: "ImageFont.FreeTypeFont", + target_width: float, + tolerance: float = 0.0, + left_context: str = "", + right_context: str = "", + casing: str = "lowercase", + known_start: str = "", + known_end: str = "", + ensure_plural: bool = False, +): + """Search English nouns (single-word) and adj+noun phrases (two-word). + + Yields results as found for SSE streaming. Phase 1 searches single nouns, + phase 2 searches two-word combinations using binary search. + """ + from unredact.pipeline.word_filter import ( + _get_nouns, + _get_nouns_plural, + ) + + nouns = _get_nouns() + plurals = _get_nouns_plural() + word_list = plurals if ensure_plural else nouns + + ks_lower = known_start.lower() + ke_lower = known_end.lower() + seen: set[str] = set() + + for word in word_list: + if ks_lower and not word.startswith(ks_lower): + continue + if ke_lower and not word.endswith(ke_lower): + continue + + display = _apply_casing(word, casing) + if display in seen: + continue + seen.add(display) + + if left_context or right_context: + full = left_context + display + right_context + full_len = font.getlength(full) + left_len = font.getlength(left_context) if left_context else 0.0 + right_len = font.getlength(right_context) if right_context else 0.0 + width = full_len - left_len - right_len + else: + width = font.getlength(display) + + error = abs(width - target_width) + if error <= tolerance: + yield SolveResult(text=display, width=float(width), error=float(error))