diff --git a/src/freshdata/enterprise/privacy_policy.py b/src/freshdata/enterprise/privacy_policy.py index 20b38e10..b7c183da 100644 --- a/src/freshdata/enterprise/privacy_policy.py +++ b/src/freshdata/enterprise/privacy_policy.py @@ -45,9 +45,11 @@ from ..adapters.polars import from_pandas, to_pandas from .config import PIIDetectionConfig from .privacy import ( + ENTITY_PATTERNS, MaskingEvent, PrivacyReport, TokenVault, + _duplicated_labels, _is_missing_scalar, _luhn_ok, detect_in_text, @@ -65,7 +67,10 @@ ) _PREVIEW_LEN = 24 -_SAMPLE_ROWS = 200 # cells sampled per column for value/entity classification +#: Classification reads every distinct non-null value of a column. Entity detection +#: runs over newline-joined chunks of at most this many characters; a single value +#: longer than this forms its own chunk. +_CLASSIFY_CHUNK_CHARS = 64_000 # ===================================================================== @@ -138,6 +143,10 @@ class PrivacyRule: (``value_regexes``), :attr:`context` keywords, and the entity/domain-pack classifier (:attr:`entity_types`). :attr:`jurisdictions` scopes the rule; an empty tuple means "any jurisdiction". + + The key for ``tokenize``/``pseudonymize`` is resolved from the rule before the + policy: :attr:`key_env` (when that variable is set), then :attr:`key`, then + the policy's ``key_env``, then the policy's ``key``. """ id: str @@ -256,10 +265,16 @@ def from_dict(cls, data: dict[str, Any]) -> CompliancePack: class PrivacyPolicy: """A complete privacy policy: jurisdiction, packs, and inline rule overrides. - Inline :attr:`rules` take priority over pack rules; within either group the - first rule that matches a column (and is in scope for :attr:`jurisdiction`) - wins. Set :attr:`minimize` to actually drop columns whose action is + Inline :attr:`rules` take priority over pack rules: any inline rule that + matches a column (and is in scope for :attr:`jurisdiction`) beats every pack + rule. Within each group the most specific classifier wins (column-name, then + entity, then value-regex, then context), and ties go to the earlier rule. + Set :attr:`minimize` to actually drop columns whose action is ``minimize`` (off by default so minimisation is an explicit opt-in). + + :attr:`key_env` and :attr:`key` are defaults for rules that set no key of + their own. A rule's ``key_env`` (when set in the environment) and ``key`` come + first, then the policy's :attr:`key_env`, then the policy's :attr:`key`. """ name: str = "privacy-policy" @@ -392,35 +407,103 @@ class _ColumnClassification: matched_by: str # which classifier kind selected the column -def _column_sample(series: pd.Series) -> tuple[list[str], str]: - values = [str(v) for v in series.dropna().tolist()[:_SAMPLE_ROWS]] - return values, "\n".join(values) +class _ColumnValues: + """Every distinct non-null value of a column, as strings, plus its text chunks. + + Values are de-duplicated in first-seen order, so a repeated value is read once. + ``chunks`` joins the values with ``"\\n"`` into pieces of at most + ``_CLASSIFY_CHUNK_CHARS`` characters, cutting only between values. + """ + + def __init__(self, series: pd.Series) -> None: + non_null = series.dropna() + try: + distinct: list[Any] = non_null.drop_duplicates().tolist() + except TypeError: # unhashable cells (lists, dicts): de-duplicate on str() + distinct = non_null.tolist() + self.values: list[str] = list(dict.fromkeys(str(v) for v in distinct)) + self.chunks: list[str] = _pack_chunks(self.values, _CLASSIFY_CHUNK_CHARS) + self._lowered: list[str] | None = None + self._luhn: bool | None = None + + @property + def lowered_chunks(self) -> list[str]: + if self._lowered is None: + self._lowered = [chunk.lower() for chunk in self.chunks] + return self._lowered + + @property + def has_luhn_candidate(self) -> bool: + if self._luhn is None: + self._luhn = any(_luhn_candidate(chunk) for chunk in self.chunks) + return self._luhn + + +def _pack_chunks(values: list[str], limit: int) -> list[str]: + """Join *values* with newlines into chunks of at most *limit* characters. + + A value longer than *limit* is never split; it becomes a chunk on its own. + """ + chunks: list[str] = [] + buf: list[str] = [] + size = 0 + for value in values: + extra = len(value) + (1 if buf else 0) + if buf and size + extra > limit: + chunks.append("\n".join(buf)) + buf, size, extra = [], 0, len(value) + buf.append(value) + size += extra + if buf: + chunks.append("\n".join(buf)) + return chunks + + +def _detect_entity_types( + column_values: _ColumnValues, column: str, cfg: PIIDetectionConfig +) -> set[str]: + """Union of the entity types :func:`detect_in_text` finds in each chunk. + + Stops once every entity type the config can report has been found, since + further chunks cannot add to the set. + """ + reachable = {e for e in cfg.entities if e in ENTITY_PATTERNS} if cfg.use_regex else set() + reachable |= {str(c.get("name", "CUSTOM")) for c in cfg.custom_patterns} + detected: set[str] = set() + for chunk in column_values.chunks: + if not chunk: + continue + detected.update(e.entity_type for e in detect_in_text(chunk, column=column, config=cfg)) + if reachable <= detected: + break + return detected def _luhn_candidate(text: str) -> bool: return any(_luhn_ok(m) for m in re.findall(r"(?:\d[ -]?){13,19}", text)) -#: Classifier specificity — a more specific signal wins when several rules match -#: the same column (column-name beats entity beats value-regex beats context). +#: Classifier specificity — when several rules of the same group (inline or pack) +#: match a column, the more specific signal wins (column-name beats entity beats +#: value-regex beats context). A matching inline rule always beats a pack rule. _CLASSIFIER_SPECIFICITY = {"column-name": 4, "entity": 3, "regex": 2, "context": 1} def _rule_matches_column( rule: PrivacyRule, column: str, - sample_values: list[str], - sample_text: str, + column_values: _ColumnValues, detected: set[str], ) -> str | None: """Return the *strongest* classifier kind that matched, or ``None``. A rule may match through several classifiers; we report the most specific so that, e.g., an exact column-name hit outranks a loose ``context`` keyword that - merely appears in another column's name. + merely appears in another column's name. The regex classifier reads each + distinct value; the Luhn and context checks read each chunk. """ col_l = column.lower() - if rule.requires_luhn and not _luhn_candidate(sample_text): + if rule.requires_luhn and not column_values.has_luhn_candidate: return None # column-name classifier if any(col_l == c.lower() for c in rule.columns) or any( @@ -432,15 +515,15 @@ def _rule_matches_column( return "entity" # regex (value) classifier if rule._compiled_vals and any( - p.search(v) for p in rule._compiled_vals for v in sample_values + p.search(v) for p in rule._compiled_vals for v in column_values.values ): return "regex" # context classifier — keywords in the *surrounding data values*. Column-name # signals are the column-name classifier's job; matching context against the # name too would let a generic word (e.g. "card") over-claim a column. if rule.context: - haystack = sample_text.lower() - if any(k.lower() in haystack for k in rule.context): + keywords = [k.lower() for k in rule.context] + if any(k in chunk for chunk in column_values.lowered_chunks for k in keywords): return "context" return None @@ -448,27 +531,66 @@ def _rule_matches_column( def classify_columns( df: Any, policy: PrivacyPolicy, *, jurisdiction: str | None = None ) -> dict[str, _ColumnClassification]: - """Classify each column under *policy* without mutating the data.""" + """Classify each column under *policy* without mutating the data. + + Every distinct non-null value of each column is read, so a sensitive value is + found wherever it sits in the column. Results are keyed by ``str(label)``, so + column labels must be unique and stay distinct once stringified (``1`` and + ``"1"`` together raise ``ValueError``). + """ frame = to_pandas(df) + _column_label_map(frame, "classify_columns") + return _classify(frame, policy, jurisdiction)[0] + + +def _column_label_map(frame: pd.DataFrame, caller: str) -> dict[str, Any]: + """Map each column's report key, ``str(label)``, back to the original label. + + Raises ``ValueError`` when a label is duplicated or when distinct labels share + a string form, since each report key must address exactly one column. + """ + duplicated = _duplicated_labels(frame) + if duplicated: + raise ValueError(f"{caller} requires unique column labels; duplicated: {duplicated}") + labels: dict[str, Any] = {} + colliding: dict[str, list[Any]] = {} + for label in frame.columns: + key = str(label) + if key in labels: + colliding.setdefault(key, [labels[key]]).append(label) + else: + labels[key] = label + if colliding: + raise ValueError( + f"{caller} requires column labels that stay distinct as strings; " + f"colliding: {list(colliding.values())}" + ) + return labels + + +def _classify( + frame: pd.DataFrame, policy: PrivacyPolicy, jurisdiction: str | None +) -> tuple[dict[str, _ColumnClassification], dict[str, int]]: + """Classifications plus the number of distinct values read per column.""" rules = policy.effective_rules(jurisdiction) + inline = {id(rule) for rule in policy.rules} cfg = policy.detection_config or PIIDetectionConfig() result: dict[str, _ColumnClassification] = {} + scanned: dict[str, int] = {} for col in frame.columns: - sample_values, sample_text = _column_sample(frame[col]) - detected: set[str] = set() - if sample_text: - detected = {e.entity_type for e in detect_in_text( - sample_text, column=str(col), config=cfg)} + column_values = _ColumnValues(frame[col]) + scanned[str(col)] = len(column_values.values) + detected = _detect_entity_types(column_values, str(col), cfg) chosen: PrivacyRule | None = None matched_by = "" - best_rank = -1 + best_rank = (False, -1) for rule in rules: # inline rules precede pack rules - kind = _rule_matches_column( - rule, str(col), sample_values, sample_text, detected) + kind = _rule_matches_column(rule, str(col), column_values, detected) if not kind: continue - rank = _CLASSIFIER_SPECIFICITY[kind] - # most specific classifier wins; ties broken by rule order (earlier wins) + # an inline rule beats any pack rule; within a group the most specific + # classifier wins; ties broken by rule order (earlier wins) + rank = (id(rule) in inline, _CLASSIFIER_SPECIFICITY[kind]) if rank > best_rank: best_rank, chosen, matched_by = rank, rule, kind if chosen is not None: @@ -484,7 +606,7 @@ def classify_columns( classification=f"detected PII: {', '.join(sorted(detected))}", risk=risk_level_for(top), matched_by="entity", ) - return result + return result, scanned # ===================================================================== @@ -493,11 +615,17 @@ def classify_columns( def _resolve_key(rule: PrivacyRule | None, policy: PrivacyPolicy) -> str | None: - for env in ((rule.key_env if rule else None), policy.key_env): - if env and os.environ.get(env): - return os.environ[env] - if rule and rule.key: - return rule.key + """Rule settings first: ``rule.key_env``, ``rule.key``, ``policy.key_env``, ``policy.key``. + + An environment variable counts only when it is set to a non-empty value. + """ + if rule is not None: + if rule.key_env and os.environ.get(rule.key_env): + return os.environ[rule.key_env] + if rule.key: + return rule.key + if policy.key_env and os.environ.get(policy.key_env): + return os.environ[policy.key_env] return policy.key @@ -585,11 +713,16 @@ def apply_privacy_policy( tokenisation requires a vault and key — supply ``vault=`` to share one, or let the policy/rule vault settings build it; a key must come from ``key``/``key_env``. Report previews are redacted unless ``audit_include_pii=True``. + + Column labels need not be strings. Report entries are keyed by ``str(label)``, + so labels must be unique and stay distinct once stringified; otherwise + ``ValueError`` is raised. """ frame = to_pandas(df).copy() + labels = _column_label_map(frame, "apply_privacy_policy") juris = Jurisdiction.coerce(jurisdiction or policy.jurisdiction) cfg = policy.detection_config or PIIDetectionConfig() - classifications = classify_columns(frame, policy, jurisdiction=juris.value) + classifications, values_scanned = _classify(frame, policy, juris.value) events: list[MaskingEvent] = [] violations: list[dict[str, Any]] = [] @@ -626,7 +759,7 @@ def apply_privacy_policy( "type": "reversible_not_allowed", "detail": f"pack {pack.name} forbids reversible tokenisation"}) - series = frame[col] + series = frame[labels[col]] n_cells = int(series.notna().sum()) reversible = False format_preserving = False @@ -661,8 +794,14 @@ def apply_privacy_policy( unprotected.append(col) elif action is Action.QUARANTINE: - new = series.where(series.isna(), _QUARANTINE_PLACEHOLDER) - frame[col] = new + try: + new = series.where(series.isna(), _QUARANTINE_PLACEHOLDER) + except (TypeError, ValueError): + # Nullable Int64/boolean and categorical columns cannot hold the + # string placeholder; substitute on object values, as an object + # column would. Missing cells stay missing. + new = series.astype(object).where(series.isna(), _QUARANTINE_PLACEHOLDER) + frame[labels[col]] = new quarantined.append(col) touched.append(col) changed_cols.append(col) @@ -709,7 +848,7 @@ def apply_privacy_policy( changed += 1 if not sample_masked: sample_masked, sample_original = masked, original - frame[col] = pd.Series(new_values, index=series.index) + frame[labels[col]] = pd.Series(new_values, index=series.index) if changed: touched.append(col) changed_cols.append(col) @@ -723,7 +862,9 @@ def apply_privacy_policy( )) if drop_cols: - frame.drop(columns=[c for c in drop_cols if c in frame.columns], inplace=True) + frame.drop( + columns=[labels[c] for c in drop_cols if labels[c] in frame.columns], inplace=True + ) detected = list(classifications.keys()) trust_dimension = { @@ -753,7 +894,11 @@ def apply_privacy_policy( cells_changed=cells_changed, columns_changed=tuple(dict.fromkeys(changed_cols + drop_cols)), events=events, - metadata={"quarantined_columns": quarantined, "dropped_columns": drop_cols}, + metadata={ + "quarantined_columns": quarantined, + "dropped_columns": drop_cols, + "classification_values_scanned": values_scanned, + }, policy_name=policy.name, jurisdiction=juris.value, compliance_pack=tuple(sorted(packs_used)), diff --git a/tests/test_privacy_missing_and_labels.py b/tests/test_privacy_missing_and_labels.py index c1dce6cf..12377f36 100644 --- a/tests/test_privacy_missing_and_labels.py +++ b/tests/test_privacy_missing_and_labels.py @@ -91,9 +91,6 @@ def test_drop_counts_only_present_cells(dtype): (dtype, action) for action in ("tokenize", "pseudonymize", "redact", "quarantine") for dtype in _NULLABLE_COLUMNS - # quarantine writes a string placeholder with Series.where, which masked - # Int64/boolean arrays reject; that is separate from missing-value handling. - if not (action == "quarantine" and dtype in ("Int64", "boolean")) ] diff --git a/tests/test_privacy_policy_classification.py b/tests/test_privacy_policy_classification.py new file mode 100644 index 00000000..614eb602 --- /dev/null +++ b/tests/test_privacy_policy_classification.py @@ -0,0 +1,380 @@ +"""Privacy policy classification: value coverage, rule priority, keys and labels. + +* #246: classification reads every distinct value, not the first 200 cells. +* #284: a matching inline rule beats every pack rule. +* #285: key precedence is rule.key_env, rule.key, policy.key_env, policy.key. +* #232 (part 5): non-string column labels are looked up by their original label. +* quarantine substitutes its placeholder on nullable Int64/boolean columns too. +""" + +from __future__ import annotations + +import math + +import pandas as pd +import pytest + +from freshdata.enterprise import ( + PIIDetectionConfig, + PrivacyPolicy, + PrivacyRule, + apply_privacy_policy, + classify_columns, + load_compliance_pack, + privacy_policy, +) + +KEY = "unit-test-key" + + +@pytest.fixture +def detect_calls(monkeypatch): + """Count the detect_in_text calls made during classification.""" + calls: list[int] = [] + real = privacy_policy.detect_in_text + + def counting(text, **kwargs): + calls.append(len(text)) + return real(text, **kwargs) + + monkeypatch.setattr(privacy_policy, "detect_in_text", counting) + return calls + + +def _distinct_words(n: int, width: int) -> list[str]: + """*n* distinct lowercase words of *width* characters that match no detector.""" + letters = "abcdefghijklmnopqrstuvwxyz" + words = [] + for i in range(n): + stem = "".join(letters[(i // 26**k) % 26] for k in range(3)) + words.append(stem + "q" * (width - len(stem))) + return words + + +# -------------------------------------------------------------------------- +# #246: classify from all distinct values +# -------------------------------------------------------------------------- + +_PII = "patient email jane.roe@clinic.org ssn 123-45-6789" + + +@pytest.mark.parametrize("n_before", [10, 250, 10_000]) +def test_issue_246_late_pii_is_classified(n_before): + pol = PrivacyPolicy(packs=(load_compliance_pack("hipaa"),), jurisdiction="US") + df = pd.DataFrame({"notes": ["follow-up scheduled"] * n_before + [_PII]}) + out, rep = apply_privacy_policy(df, pol) + assert out["notes"].iloc[-1] == "patient email ssn " + assert sorted(rep.classifications) == ["notes"] + assert rep.metadata["classification_values_scanned"] == {"notes": 2} + + +def test_late_value_regex_match_is_classified(): + rule = PrivacyRule(id="ticket", action="redact", value_regexes=(r"\bTKT-\d{6}\b",)) + df = pd.DataFrame({"ref": [f"note {i}" for i in range(5_000)] + ["TKT-123456"]}) + cls = classify_columns(df, PrivacyPolicy(rules=(rule,))) + assert cls["ref"].rule is rule + assert cls["ref"].matched_by == "regex" + + +def test_late_context_keyword_is_classified(): + rule = PrivacyRule(id="conf", action="quarantine", context=("confidential",)) + df = pd.DataFrame({"memo": [f"memo {i}" for i in range(5_000)] + ["CONFIDENTIAL draft"]}) + out, rep = apply_privacy_policy(df, PrivacyPolicy(rules=(rule,))) + assert rep.classifications["memo"]["matched_by"] == "context" + assert set(out["memo"]) == {""} + + +def test_late_luhn_value_satisfies_requires_luhn(): + rule = PrivacyRule(id="pan", action="redact", requires_luhn=True, value_regexes=(r"\d{4}",)) + head = ["order 1234 5678 9012 3456"] * 300 # 16 digits, fails Luhn + policy = PrivacyPolicy(rules=(rule,)) + assert "ref" not in classify_columns(pd.DataFrame({"ref": head}), policy) + cls = classify_columns(pd.DataFrame({"ref": head + ["4111 1111 1111 1111"]}), policy) + assert cls["ref"].rule is rule + + +def test_detect_calls_match_chunk_count(detect_calls): + values = _distinct_words(640, 999) + chars = len("\n".join(values)) + cls = classify_columns(pd.DataFrame({"words": values}), PrivacyPolicy()) + assert cls == {} + assert len(detect_calls) == math.ceil(chars / privacy_policy._CLASSIFY_CHUNK_CHARS) + assert max(detect_calls) <= privacy_policy._CLASSIFY_CHUNK_CHARS + + +def test_repeated_values_are_read_once(detect_calls): + values = _distinct_words(640, 999) + classify_columns(pd.DataFrame({"words": values * 5}), PrivacyPolicy()) + assert len(detect_calls) == 10 + detect_calls.clear() + _, rep = apply_privacy_policy(pd.DataFrame({"same": ["repeated"] * 10_000}), PrivacyPolicy()) + assert len(detect_calls) == 1 + assert rep.metadata["classification_values_scanned"] == {"same": 1} + + +def test_detection_stops_once_every_entity_type_is_found(detect_calls): + values = ["mail a@b.com", *_distinct_words(200, 999)] + policy = PrivacyPolicy(detection_config=PIIDetectionConfig(entities=("EMAIL",))) + cls = classify_columns(pd.DataFrame({"notes": values}), policy) + assert cls["notes"].classification == "detected PII: EMAIL" + assert len(detect_calls) == 1 + + +def test_oversized_value_forms_its_own_chunk(): + limit = privacy_policy._CLASSIFY_CHUNK_CHARS + big = "q" * (limit + 10) + chunks = privacy_policy._pack_chunks(["a", big, "b", "c"], limit) + assert chunks == ["a", big, "b\nc"] + assert privacy_policy._pack_chunks([], limit) == [] + + +def test_values_scanned_ignores_missing_cells(): + df = pd.DataFrame({"s": pd.array(["x", "x", None, "y"], dtype="string"), "n": [1, 1, 2, 2]}) + _, rep = apply_privacy_policy(df, PrivacyPolicy()) + assert rep.metadata["classification_values_scanned"] == {"s": 2, "n": 2} + + +def test_unhashable_cells_are_classified(): + df = pd.DataFrame({"tags": [["a"], ["a"], ["mail a@b.com"]]}) + _, rep = apply_privacy_policy(df, PrivacyPolicy()) + assert rep.classifications["tags"]["classification"] == "detected PII: EMAIL" + assert rep.metadata["classification_values_scanned"] == {"tags": 2} + + +# -------------------------------------------------------------------------- +# #284: inline rules take priority over pack rules +# -------------------------------------------------------------------------- + + +def _gdpr_policy(*rules: PrivacyRule, jurisdiction: str = "EU") -> PrivacyPolicy: + return PrivacyPolicy( + rules=rules, packs=(load_compliance_pack("gdpr"),), jurisdiction=jurisdiction + ) + + +def test_issue_284_inline_entity_rule_beats_pack_column_name_rule(): + inline = PrivacyRule(id="inline.email_drop", action="drop", entity_types=("EMAIL",)) + out, rep = apply_privacy_policy( + pd.DataFrame({"email": ["a@b.com"], "v": [1]}), _gdpr_policy(inline) + ) + assert rep.classifications["email"]["rule_id"] == "inline.email_drop" + assert rep.classifications["email"]["matched_by"] == "entity" + assert list(out.columns) == ["v"] + + +def test_inline_context_rule_beats_pack_column_name_rule(): + inline = PrivacyRule(id="inline.hold", action="quarantine", context=("contact",)) + df = pd.DataFrame({"email": ["contact a@b.com", "c@d.com"]}) + out, rep = apply_privacy_policy(df, _gdpr_policy(inline)) + assert rep.classifications["email"]["rule_id"] == "inline.hold" + assert rep.classifications["email"]["matched_by"] == "context" + assert list(out["email"]) == ["", ""] + + +def test_specificity_still_applies_within_inline_rules(): + by_context = PrivacyRule(id="by_context", action="redact", context=("contact",)) + by_name = PrivacyRule(id="by_name", action="drop", columns=("email",)) + by_name_too = PrivacyRule(id="by_name_too", action="quarantine", columns=("email",)) + df = pd.DataFrame({"email": ["contact a@b.com"]}) + cls = classify_columns(df, _gdpr_policy(by_context, by_name, by_name_too)) + assert cls["email"].rule is by_name + assert cls["email"].matched_by == "column-name" + + +def test_out_of_scope_inline_rule_does_not_override_pack(): + us_only = PrivacyRule( + id="us.email_drop", action="drop", entity_types=("EMAIL",), jurisdictions=("US",) + ) + out, rep = apply_privacy_policy(pd.DataFrame({"email": ["a@b.com"]}), _gdpr_policy(us_only)) + assert rep.classifications["email"]["rule_id"] == "gdpr.email" + assert list(out.columns) == ["email"] + + +def test_pack_only_policy_is_unchanged(): + df = pd.DataFrame( + { + "email": ["a@b.com"], + "phone_number": ["+44 20 7946 0958"], + "notes": ["reach me at c@d.com"], + } + ) + cls = classify_columns(df, _gdpr_policy()) + assert {c: (v.rule.id, v.matched_by) for c, v in cls.items() if v.rule} == { + "email": ("gdpr.email", "column-name"), + "phone_number": ("gdpr.phone", "column-name"), + "notes": ("gdpr.email", "entity"), + } + + +# -------------------------------------------------------------------------- +# #285: key precedence +# -------------------------------------------------------------------------- + +_RULE_ENV = "FRESHDATA_TEST_RULE_KEY" +_POLICY_ENV = "FRESHDATA_TEST_POLICY_KEY" + + +@pytest.fixture +def clean_key_env(monkeypatch): + for name in (_RULE_ENV, _POLICY_ENV, "ORG_DEFAULT_KEY"): + monkeypatch.delenv(name, raising=False) + return monkeypatch + + +def _tokenize(policy: PrivacyPolicy) -> str: + df = pd.DataFrame({"ssn": ["123-45-6789"]}) + return apply_privacy_policy(df, policy)[0]["ssn"][0] + + +def test_issue_285_rule_key_beats_policy_key_env(clean_key_env): + df = pd.DataFrame({"ssn": ["123-45-6789"]}) + rule = PrivacyRule(id="t", action="tokenize", columns=("ssn",), key="rule-specific-key") + pol = PrivacyPolicy(rules=(rule,), key_env="ORG_DEFAULT_KEY") + before = apply_privacy_policy(df, pol)[0]["ssn"][0] + clean_key_env.setenv("ORG_DEFAULT_KEY", "org-wide-key") + after = apply_privacy_policy(df, pol)[0]["ssn"][0] + assert before == after == "tok_05817a27d410ecc4" + + +def test_rule_key_env_beats_rule_key(clean_key_env): + rule = PrivacyRule(id="t", action="tokenize", key="rule-literal", key_env=_RULE_ENV) + pol = PrivacyPolicy(rules=(rule,), key="policy-literal", key_env=_POLICY_ENV) + clean_key_env.setenv(_POLICY_ENV, "policy-env") + assert privacy_policy._resolve_key(rule, pol) == "rule-literal" + clean_key_env.setenv(_RULE_ENV, "rule-env") + assert privacy_policy._resolve_key(rule, pol) == "rule-env" + clean_key_env.setenv(_RULE_ENV, "") + assert privacy_policy._resolve_key(rule, pol) == "rule-literal" + + +def test_rule_key_env_without_rule_key_falls_back_to_policy(clean_key_env): + rule = PrivacyRule(id="t", action="tokenize", key_env=_RULE_ENV) + pol = PrivacyPolicy(rules=(rule,), key="policy-literal", key_env=_POLICY_ENV) + assert privacy_policy._resolve_key(rule, pol) == "policy-literal" + clean_key_env.setenv(_POLICY_ENV, "policy-env") + assert privacy_policy._resolve_key(rule, pol) == "policy-env" + + +@pytest.mark.parametrize("rule", [None, PrivacyRule(id="keyless", action="tokenize")]) +def test_keyless_rule_uses_policy_key_env_then_policy_key(clean_key_env, rule): + pol = PrivacyPolicy(key="policy-literal", key_env=_POLICY_ENV) + assert privacy_policy._resolve_key(rule, pol) == "policy-literal" + clean_key_env.setenv(_POLICY_ENV, "policy-env") + assert privacy_policy._resolve_key(rule, pol) == "policy-env" + assert privacy_policy._resolve_key(rule, PrivacyPolicy()) is None + + +def test_keyless_rule_tokens_follow_policy_key(clean_key_env): + keyless = PrivacyRule(id="t", action="tokenize", columns=("ssn",)) + keyed = PrivacyRule(id="t", action="tokenize", columns=("ssn",), key="policy-env") + clean_key_env.setenv(_POLICY_ENV, "policy-env") + via_policy = _tokenize(PrivacyPolicy(rules=(keyless,), key_env=_POLICY_ENV)) + assert via_policy == _tokenize(PrivacyPolicy(rules=(keyed,))) + assert via_policy.startswith("tok_") + + +def test_tokenize_without_any_key_still_raises(clean_key_env): + rule = PrivacyRule(id="t", action="tokenize", columns=("ssn",), key_env=_RULE_ENV) + pol = PrivacyPolicy(rules=(rule,), key_env=_POLICY_ENV) + with pytest.raises(ValueError, match="tokenize requires a key"): + apply_privacy_policy(pd.DataFrame({"ssn": ["123-45-6789"]}), pol) + + +# -------------------------------------------------------------------------- +# #232 part 5: non-string column labels +# -------------------------------------------------------------------------- + + +def test_issue_232_integer_labels_are_masked(): + pol = PrivacyPolicy(packs=(load_compliance_pack("hipaa"),), jurisdiction="US") + df = pd.DataFrame({0: ["a@b.com", "c@d.com"], 1: [1, 2]}) + out, rep = apply_privacy_policy(df, pol) + assert list(out[0]) == ["", ""] + assert list(out.columns) == [0, 1] + assert list(out[1]) == [1, 2] + assert sorted(rep.classifications) == ["0"] + assert rep.columns_changed == ("0",) + assert list(df[0]) == ["a@b.com", "c@d.com"] + + +def test_classify_columns_reports_string_keys_for_integer_labels(): + pol = PrivacyPolicy(packs=(load_compliance_pack("hipaa"),), jurisdiction="US") + cls = classify_columns(pd.DataFrame({0: ["a@b.com"], 1: [1]}), pol) + assert list(cls) == ["0"] + assert cls["0"].column == "0" + + +def test_drop_with_integer_labels(): + rule = PrivacyRule(id="d", action="drop", columns=("0",)) + df = pd.DataFrame({0: ["x", None], 1: ["y", "z"]}) + out, rep = apply_privacy_policy(df, PrivacyPolicy(rules=(rule,))) + assert list(out.columns) == [1] + assert rep.metadata["dropped_columns"] == ["0"] + assert rep.cells_changed == 1 + + +def test_quarantine_with_integer_labels(): + rule = PrivacyRule(id="q", action="quarantine", columns=("1",)) + df = pd.DataFrame({0: ["x", "y"], 1: ["s", None]}) + out, rep = apply_privacy_policy(df, PrivacyPolicy(rules=(rule,))) + assert list(out.columns) == [0, 1] + assert out[1].iloc[0] == "" + assert out[1].iloc[1] is None + assert rep.metadata["quarantined_columns"] == ["1"] + + +def test_tokenize_with_integer_labels_matches_string_labels(): + rule = PrivacyRule(id="t", action="tokenize", columns=("7",), key=KEY) + policy = PrivacyPolicy(rules=(rule,)) + out_int, _ = apply_privacy_policy(pd.DataFrame({7: ["123-45-6789"]}), policy) + out_str, _ = apply_privacy_policy(pd.DataFrame({"7": ["123-45-6789"]}), policy) + assert list(out_int[7]) == list(out_str["7"]) + + +@pytest.mark.parametrize("call", [classify_columns, apply_privacy_policy]) +def test_labels_colliding_as_strings_raise(call): + df = pd.DataFrame({1: ["a@b.com"], "1": ["c@d.com"]}) + with pytest.raises(ValueError, match=r"distinct as strings; colliding: \[\[1, '1'\]\]"): + call(df, PrivacyPolicy()) + + +@pytest.mark.parametrize("call", [classify_columns, apply_privacy_policy]) +def test_duplicate_labels_raise(call): + df = pd.DataFrame([["a@b.com", "c@d.com"]], columns=["email", "email"]) + with pytest.raises(ValueError, match=r"unique column labels; duplicated: \['email'\]"): + call(df, PrivacyPolicy()) + + +# -------------------------------------------------------------------------- +# quarantine on nullable and categorical dtypes +# -------------------------------------------------------------------------- + +_QUARANTINE_COLUMNS = { + "Int64": lambda: pd.array([1, pd.NA, 3], dtype="Int64"), + "boolean": lambda: pd.array([True, pd.NA, False], dtype="boolean"), + "string": lambda: pd.array(["a", pd.NA, "b"], dtype="string"), + "category": lambda: pd.Series(["a", None, "b"], dtype="category"), + "object_pd_NA": lambda: pd.Series(["a", pd.NA, "b"], dtype=object), +} + + +@pytest.mark.parametrize("dtype", list(_QUARANTINE_COLUMNS)) +def test_quarantine_nullable_columns(dtype): + rule = PrivacyRule(id="q", action="quarantine", columns=("c",)) + df = pd.DataFrame({"c": _QUARANTINE_COLUMNS[dtype](), "keep": [1, 2, 3]}) + before = df.copy() + out, rep = apply_privacy_policy(df, PrivacyPolicy(rules=(rule,))) + assert out["c"].iloc[0] == out["c"].iloc[2] == "" + assert pd.isna(out["c"].iloc[1]) + assert list(out["keep"]) == [1, 2, 3] + assert rep.cells_changed == 2 + assert rep.metadata["quarantined_columns"] == ["c"] + pd.testing.assert_frame_equal(df, before) + + +def test_quarantine_string_dtype_is_unchanged(): + rule = PrivacyRule(id="q", action="quarantine", columns=("c",)) + df = pd.DataFrame({"c": pd.array(["a", pd.NA], dtype="string")}) + out, _ = apply_privacy_policy(df, PrivacyPolicy(rules=(rule,))) + assert str(out["c"].dtype) == "string" + assert out["c"].iloc[1] is pd.NA