Merge pull request #165 from DataFog/fix/redact-token-document-order

fix: number redaction tokens in document order

Author: sidmohan0

CHANGELOG.MD

@@ -1,5 +1,41 @@
 # ChangeLog
 
+## [Unreleased]
+
+### `datafog-python`
+
+#### Fixed
+
+- **Redaction token numbering now follows document order**: `redact()`
+  previously assigned per-type counters while replacing spans right-to-left,
+  so with multiple entities of the same type the *last* occurrence received
+  `_1` (two emails redacted as `[EMAIL_2] ... [EMAIL_1]`). Tokens for the
+  `token` and `pseudonymize` strategies are now numbered left-to-right, so
+  the first occurrence is always `_1`. This changes redacted output text for
+  multi-entity inputs; `RedactResult.entities` is now also ordered by
+  document position (ascending) instead of descending. Replacement values
+  in `mapping` are now always read from the original input text, fixing a
+  case where an already-inserted token could leak into a mapping value when
+  entity spans overlapped.
+- **`redact()` now resolves overlapping and duplicate entity spans**:
+  `redact()` is public API and accepts arbitrary entity lists, but applied
+  every span verbatim — overlapping spans (e.g. a phone number and a
+  sub-span of it) corrupted the output with stray fragments like
+  `[PHONE_1]]`. For each group of overlapping spans only one is now
+  redacted: the longest, breaking ties by entity type priority and then
+  confidence — the same suppression rule the regex scan pipeline already
+  used. Suppressed spans are omitted from `RedactResult.entities` and
+  `mapping`. Output for non-overlapping entity lists (including everything
+  produced by `scan()`) is unchanged.
+- **`mask` strategy mapping no longer collides on same-length values**:
+  the mask replacement (`***…`) previously keyed `mapping` directly, so two
+  distinct values of the same length produced the same key and only one
+  original survived. Mask mappings are now keyed by per-type indexed keys
+  in document order (`[EMAIL_MASK_1]`, `[EMAIL_MASK_2]`), preserving every
+  original value. Consumers that looked up mask mappings by the mask string
+  must switch to the indexed keys; `token`, `hash`, and `pseudonymize`
+  mapping keys are unchanged.
+
 ## [2026-07-02]
 
 ### `datafog-python` [4.7.0]

datafog/engine.py

@@ -173,6 +173,7 @@ def _suppress_overlapping_entities(entities: list[Entity]) -> list[Entity]:
         key=lambda item: (
             -_entity_length(item),
             -ENTITY_TYPE_PRIORITY.get(item.type, 0),
+            -item.confidence,
             item.start,
             item.end,
             item.type,
@@ -465,7 +466,18 @@ def redact(
     entities: list[Entity],
     strategy: str = "token",
 ) -> RedactResult:
-    """Redact PII entities from text."""
+    """Redact PII entities from text.
+
+    Callers may pass arbitrary entity lists (not just ``scan()`` output), so
+    overlapping or duplicate spans are resolved here: for each overlapping
+    group only one span is redacted — the longest, breaking ties by entity
+    type priority and then confidence. Suppressed spans are omitted from
+    ``RedactResult.entities`` and ``mapping``.
+
+    ``mapping`` for the ``mask`` strategy is keyed by per-type indexed keys
+    (``[EMAIL_MASK_1]``) rather than the mask string, because distinct
+    same-length values produce identical masks and would collide.
+    """
     if not isinstance(text, str):
         raise TypeError("text must be a string")
     if strategy not in {"token", "mask", "hash", "pseudonymize"}:
@@ -481,14 +493,22 @@ def redact(
         for entity in entities
         if 0 <= entity.start < entity.end <= len(text) and entity.text
     ]
-    valid_entities = sorted(
-        valid_entities, key=lambda e: (e.start, e.end), reverse=True
-    )
+    # Returns non-overlapping spans sorted by document position.
+    valid_entities = _suppress_overlapping_entities(valid_entities)
 
+    # Assign replacements in document order so the per-type counters used by
+    # the token/pseudonymize strategies number the first occurrence _1; spans
+    # are applied right-to-left below so earlier offsets stay valid.
+    replacements: list[tuple[Entity, str]] = []
     for entity in valid_entities:
-        original = redacted_text[entity.start : entity.end]
+        original = text[entity.start : entity.end]
+        mapping_key: Optional[str] = None
         if strategy == "mask":
             replacement = "*" * max(len(original), 1)
+            # Distinct same-length values share a mask, so the mask string
+            # cannot key the mapping without collisions.
+            counters[entity.type] = counters.get(entity.type, 0) + 1
+            mapping_key = f"[{entity.type}_MASK_{counters[entity.type]}]"
         elif strategy == "hash":
             digest = hashlib.sha256(original.encode("utf-8")).hexdigest()[:12]
             replacement = f"[{entity.type}_{digest}]"
@@ -504,10 +524,13 @@ def redact(
             counters[entity.type] = counters.get(entity.type, 0) + 1
             replacement = f"[{entity.type}_{counters[entity.type]}]"
 
+        replacements.append((entity, replacement))
+        mapping[mapping_key if mapping_key is not None else replacement] = original
+
+    for entity, replacement in reversed(replacements):
         redacted_text = (
             redacted_text[: entity.start] + replacement + redacted_text[entity.end :]
         )
-        mapping[replacement] = original
 
     return RedactResult(
         redacted_text=redacted_text,

tests/test_engine_api.py

@@ -56,6 +56,143 @@ def test_redact_strategies(strategy: str) -> None:
     assert result.mapping
 
 
+def _entity_at(text: str, value: str, entity_type: str = "EMAIL") -> Entity:
+    start = text.index(value)
+    return Entity(
+        type=entity_type,
+        text=value,
+        start=start,
+        end=start + len(value),
+        confidence=1.0,
+        engine="regex",
+    )
+
+
+def test_redact_token_numbering_follows_document_order() -> None:
+    text = "first alpha@example.com then beta@example.com end"
+    entities = [
+        _entity_at(text, "alpha@example.com"),
+        _entity_at(text, "beta@example.com"),
+    ]
+
+    result = redact(text=text, entities=entities, strategy="token")
+
+    assert result.redacted_text == "first [EMAIL_1] then [EMAIL_2] end"
+    assert result.mapping == {
+        "[EMAIL_1]": "alpha@example.com",
+        "[EMAIL_2]": "beta@example.com",
+    }
+
+
+def test_redact_pseudonymize_numbering_follows_document_order() -> None:
+    text = "first alpha@example.com then beta@example.com end"
+    entities = [
+        _entity_at(text, "alpha@example.com"),
+        _entity_at(text, "beta@example.com"),
+    ]
+
+    result = redact(text=text, entities=entities, strategy="pseudonymize")
+
+    assert result.redacted_text == "first [EMAIL_PSEUDO_1] then [EMAIL_PSEUDO_2] end"
+    assert result.mapping == {
+        "[EMAIL_PSEUDO_1]": "alpha@example.com",
+        "[EMAIL_PSEUDO_2]": "beta@example.com",
+    }
+
+
+def _span(
+    text: str,
+    start: int,
+    end: int,
+    entity_type: str = "PHONE",
+    confidence: float = 1.0,
+    engine: str = "regex",
+) -> Entity:
+    return Entity(
+        type=entity_type,
+        text=text[start:end],
+        start=start,
+        end=end,
+        confidence=confidence,
+        engine=engine,
+    )
+
+
+@pytest.mark.parametrize("strategy", ["token", "mask", "hash", "pseudonymize"])
+def test_redact_overlapping_spans_match_longest_only(strategy: str) -> None:
+    text = "call 555-000-1111 now"
+    full = _span(text, 5, 17)
+    suffix = _span(text, 9, 17)
+
+    result = redact(text=text, entities=[full, suffix], strategy=strategy)
+    expected = redact(text=text, entities=[full], strategy=strategy)
+
+    assert result.redacted_text == expected.redacted_text
+    assert result.mapping == expected.mapping
+    assert result.entities == [full]
+
+
+@pytest.mark.parametrize("strategy", ["token", "mask", "hash", "pseudonymize"])
+def test_redact_nested_spans_keep_outer(strategy: str) -> None:
+    text = "mail alice@example.com today"
+    outer = _span(text, 5, 22, entity_type="EMAIL")
+    inner = _span(text, 11, 18, entity_type="EMAIL")
+
+    result = redact(text=text, entities=[inner, outer], strategy=strategy)
+    expected = redact(text=text, entities=[outer], strategy=strategy)
+
+    assert result.redacted_text == expected.redacted_text
+    assert result.mapping == expected.mapping
+    assert result.entities == [outer]
+
+
+def test_redact_overlapping_spans_produce_clean_token_output() -> None:
+    text = "call 555-000-1111 now"
+    entities = [_span(text, 5, 17), _span(text, 9, 17)]
+
+    result = redact(text=text, entities=entities, strategy="token")
+
+    assert result.redacted_text == "call [PHONE_1] now"
+    assert result.mapping == {"[PHONE_1]": text[5:17]}
+
+
+def test_redact_duplicate_spans_applied_once() -> None:
+    text = "mail alice@example.com today"
+    entity = _entity_at(text, "alice@example.com")
+
+    result = redact(text=text, entities=[entity, entity], strategy="token")
+
+    assert result.redacted_text == "mail [EMAIL_1] today"
+    assert result.entities == [entity]
+
+
+def test_redact_overlap_tiebreak_prefers_higher_confidence() -> None:
+    text = "Acme Corporation announced"
+    org = _span(text, 0, 16, entity_type="ORGANIZATION", confidence=0.7, engine="spacy")
+    person = _span(text, 0, 16, entity_type="PERSON", confidence=0.9, engine="gliner")
+
+    result = redact(text=text, entities=[org, person], strategy="token")
+
+    assert result.redacted_text == "[PERSON_1] announced"
+    assert result.entities == [person]
+
+
+def test_redact_mask_mapping_distinguishes_same_length_values() -> None:
+    text = "a alice@example.com b bobby@example.com c"
+    entities = [
+        _entity_at(text, "alice@example.com"),
+        _entity_at(text, "bobby@example.com"),
+    ]
+
+    result = redact(text=text, entities=entities, strategy="mask")
+
+    assert result.redacted_text == "a " + "*" * 17 + " b " + "*" * 17 + " c"
+    assert result.mapping == {
+        "[EMAIL_MASK_1]": "alice@example.com",
+        "[EMAIL_MASK_2]": "bobby@example.com",
+    }
+
+
 def test_redact_invalid_strategy_raises_value_error() -> None:
     with pytest.raises(ValueError, match="strategy must be one of"):
         redact("test", entities=[], strategy="invalid")