redac/redact.py

"""Reversible redaction.

Replace each detected PII span with a stable placeholder like [PERSON_1].
The same underlying value always maps to the same placeholder, so a document
stays internally consistent ("John" referenced twice -> [PERSON_1] twice).

The mapping placeholder -> original value is kept locally so a downstream
model's answer can be rehydrated without the model ever seeing real PII.
"""

from __future__ import annotations

import re
from typing import Dict, List, Tuple

from .detect import Entity


def _slug(label: str) -> str:
    return re.sub(r"[^A-Z0-9]+", "_", label.upper()).strip("_")


def redact(text: str, entities: List[Entity]) -> Tuple[str, Dict[str, str]]:
    """Return (redacted_text, mapping) where mapping is placeholder -> original.

    Placeholders are assigned per (label, normalized-value) so repeats reuse
    the same token. Spans are replaced right-to-left to keep offsets valid.
    """
    mapping: Dict[str, str] = {}
    value_to_token: Dict[Tuple[str, str], str] = {}
    counters: Dict[str, int] = {}

    # Replace from the end so earlier offsets stay valid.
    redacted = text
    for ent in sorted(entities, key=lambda e: e.start, reverse=True):
        key = (ent.label, ent.text.strip().lower())
        token = value_to_token.get(key)
        if token is None:
            slug = _slug(ent.label)
            counters[slug] = counters.get(slug, 0) + 1
            token = f"[{slug}_{counters[slug]}]"
            value_to_token[key] = token
            mapping[token] = ent.text
        redacted = redacted[: ent.start] + token + redacted[ent.end :]

    return redacted, mapping


def rehydrate(text: str, mapping: Dict[str, str]) -> str:
    """Swap placeholders back to their original values (longest token first to
    avoid partial overlaps like [X_1] inside [X_10])."""
    out = text
    for token in sorted(mapping, key=len, reverse=True):
        out = out.replace(token, mapping[token])
    return out