diff --git a/tests/test_name_dictionary.py b/tests/test_name_dictionary.py new file mode 100644 index 0000000..f2165b1 --- /dev/null +++ b/tests/test_name_dictionary.py @@ -0,0 +1,128 @@ +"""Tests for dictionary-based single-name solving.""" + +from unittest.mock import MagicMock, patch + +from unredact.pipeline.dictionary import solve_name_dictionary +from unredact.pipeline.solver import SolveResult + + +def _mock_font(width_map: dict[str, float]) -> MagicMock: + """Create a mock font that returns widths from a map.""" + font = MagicMock() + font.getlength.side_effect = lambda text: width_map.get(text, len(text) * 7.0) + return font + + +class TestSolveNameDictionary: + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_basic_lowercase(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["john", "jane"] + mock_lasts.return_value = ["doe", "smith"] + + font = _mock_font({"john": 30.0, "jane": 28.0, "doe": 20.0, "smith": 32.0}) + + results = solve_name_dictionary(font, 30.0, 1.0) + texts = [r.text for r in results] + assert "john" in texts + assert "doe" not in texts # 20.0 is outside tolerance + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_uppercase_casing(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["john"] + mock_lasts.return_value = [] + + font = _mock_font({"JOHN": 35.0}) + + results = solve_name_dictionary(font, 35.0, 1.0, casing="uppercase") + assert len(results) == 1 + assert results[0].text == "JOHN" + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_capitalized_casing(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["john"] + mock_lasts.return_value = [] + + font = _mock_font({"John": 32.0}) + + results = solve_name_dictionary(font, 32.0, 1.0, casing="capitalized") + assert len(results) == 1 + assert results[0].text == "John" + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_known_start_filters_and_strips(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["joe", "john", "bob"] + mock_lasts.return_value = [] + + # Gap width is for "oe" (the unknown part after "j") + font = _mock_font({"oe": 14.0, "ohn": 21.0}) + + results = solve_name_dictionary( + font, 14.0, 1.0, known_start="j", + ) + texts = [r.text for r in results] + assert "joe" in texts # "j" + "oe", "oe" width matches + assert "bob" not in texts # doesn't start with "j" + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_known_end_filters_and_strips(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["johnson", "jackson"] + mock_lasts.return_value = [] + + # Gap width is for "john" (the unknown part before "son") + # With known_end="son", right kerning context is "s" + # Width = getlength("john" + "s") - getlength("s") = 37.0 - 7.0 = 30.0 + font = _mock_font({"johns": 37.0, "s": 7.0}) + + results = solve_name_dictionary( + font, 30.0, 1.0, known_end="son", + ) + texts = [r.text for r in results] + assert "johnson" in texts + assert "jackson" not in texts # doesn't end with "son" + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_dedup_first_and_last(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["lee"] + mock_lasts.return_value = ["lee"] # same name in both lists + + font = _mock_font({"lee": 20.0}) + + results = solve_name_dictionary(font, 20.0, 1.0) + assert len(results) == 1 # no duplicates + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_sorted_by_error(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["john", "jane"] + mock_lasts.return_value = [] + + font = _mock_font({"john": 30.5, "jane": 30.0}) + + results = solve_name_dictionary(font, 30.0, 1.0) + assert len(results) == 2 + assert results[0].error <= results[1].error + + @patch("unredact.pipeline.word_filter._get_associate_firsts") + @patch("unredact.pipeline.word_filter._get_associate_lasts") + def test_context_chars(self, mock_lasts, mock_firsts): + mock_firsts.return_value = ["john"] + mock_lasts.return_value = [] + + font = _mock_font({ + "": 45.0, + "<": 5.0, + ">": 5.0, + }) + + results = solve_name_dictionary( + font, 30.0, 5.0, left_context="<", right_context=">", + ) + # Width = 45 - 5 - 5 = 35, error = |35-30| = 5, within tolerance + assert len(results) == 1 + assert results[0].text == "john" diff --git a/unredact/pipeline/dictionary.py b/unredact/pipeline/dictionary.py index 9256b50..3bb0e4c 100644 --- a/unredact/pipeline/dictionary.py +++ b/unredact/pipeline/dictionary.py @@ -126,3 +126,99 @@ def solve_full_name_dictionary( results.sort(key=lambda r: (r.error, r.text)) return results + + +def solve_name_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 = "", +) -> list[SolveResult]: + """Match single associate names against target width. + + Loads first and last name lists, applies casing, filters by + known_start/known_end, and measures the unknown portion against + the target width. + """ + from unredact.pipeline.word_filter import ( + _get_associate_firsts, + _get_associate_lasts, + ) + + firsts = _get_associate_firsts() + lasts = _get_associate_lasts() + + # Combine and dedup + seen: set[str] = set() + names: list[str] = [] + for name in firsts + lasts: + if name not in seen: + seen.add(name) + names.append(name) + + results: list[SolveResult] = [] + seen_results: set[str] = set() + + ks_lower = known_start.lower() + ke_lower = known_end.lower() + + for name in names: + # Filter by known start/end (case-insensitive on raw lowercase name) + if ks_lower and not name.startswith(ks_lower): + continue + if ke_lower and not name.endswith(ke_lower): + continue + + # Apply casing to the full name (for display) + if casing == "uppercase": + display = name.upper() + elif casing == "capitalized": + display = name.title() + else: + display = name + + if display in seen_results: + continue + seen_results.add(display) + + # The unknown portion is the name minus known_start and known_end + end_idx = len(name) - len(known_end) if known_end else len(name) + unknown = name[len(known_start):end_idx] + + if not unknown: + continue # entire name is known, nothing to measure + + # Apply same casing to the unknown portion + if casing == "uppercase": + unknown_display = unknown.upper() + elif casing == "capitalized": + # Mid-word after known_start: lowercase. At word start: title case. + unknown_display = unknown.lower() if known_start else unknown.title() + else: + unknown_display = unknown + + # Determine kerning context + # If known_start is set, its last char is left context for the unknown part + effective_left = known_start[-1] if known_start else left_context + effective_right = known_end[0] if known_end else right_context + + # Measure width of unknown portion with kerning context + if effective_left or effective_right: + full = effective_left + unknown_display + effective_right + full_len = font.getlength(full) + left_len = font.getlength(effective_left) if effective_left else 0.0 + right_len = font.getlength(effective_right) if effective_right else 0.0 + width = full_len - left_len - right_len + else: + width = font.getlength(unknown_display) + + error = abs(width - target_width) + if error <= tolerance: + results.append(SolveResult(text=display, width=float(width), error=float(error))) + + results.sort(key=lambda r: (r.error, r.text)) + return results