audit-labs/audit-tools

  1"""Validation and population preparation."""
  2
  3from __future__ import annotations
  4
  5from dataclasses import dataclass
  6
  7import pandas as pd
  8
  9from .io import AuditSamplingError
 10from .methods import parse_key_floats, parse_key_ints
 11
 12
 13@dataclass
 14class ValidationResult:
 15    population: pd.DataFrame
 16    excluded_rows: pd.DataFrame
 17    duplicate_rows: pd.DataFrame
 18    effective_id_column: str
 19    id_column_omitted: bool
 20    blank_id_count: int
 21    duplicate_id_count: int
 22    warnings: list[str]
 23    strata_counts: dict[str, int]
 24    strata_proportions: dict[str, float]
 25
 26
 27def blank_id_mask(series: pd.Series) -> pd.Series:
 28    return series.isna() | (series.astype("string").fillna("").str.strip() == "")
 29
 30
 31def validate_and_prepare(population: pd.DataFrame, options) -> ValidationResult:
 32    warnings: list[str] = []
 33    excluded_parts: list[pd.DataFrame] = []
 34    working = population.copy()
 35
 36    if options.id_column:
 37        if options.id_column not in working.columns:
 38            raise AuditSamplingError(f"ID column not found: {options.id_column}")
 39        effective_id_column = options.id_column
 40        id_column_omitted = False
 41    else:
 42        effective_id_column = "_audit_row_id"
 43        id_column_omitted = True
 44        working[effective_id_column] = working["_source_row_number"]
 45        warnings.append(
 46            "--id-column omitted; using _source_row_number as generated _audit_row_id."
 47        )
 48
 49    if options.stratify_column and options.stratify_column not in working.columns:
 50        raise AuditSamplingError(
 51            f"Stratification column not found: {options.stratify_column}"
 52        )
 53
 54    id_blank_mask = blank_id_mask(working[effective_id_column])
 55    blank_id_count = int(id_blank_mask.sum())
 56    if blank_id_count:
 57        if options.exclude_blank_id:
 58            excluded = working.loc[id_blank_mask].copy()
 59            excluded["_exclusion_reason"] = "Blank ID"
 60            excluded_parts.append(excluded)
 61            working = working.loc[~id_blank_mask].copy()
 62        else:
 63            warnings.append(
 64                f"{blank_id_count} row(s) have blank IDs and were retained."
 65            )
 66
 67    nonblank_ids = ~blank_id_mask(working[effective_id_column])
 68    duplicate_mask = working.loc[nonblank_ids, effective_id_column].duplicated(
 69        keep=False
 70    )
 71    duplicate_rows = working.loc[nonblank_ids].loc[duplicate_mask].copy()
 72    duplicate_id_count = len(duplicate_rows)
 73    if duplicate_id_count:
 74        if options.dedupe_id == "fail":
 75            raise AuditSamplingError(
 76                f"Duplicate IDs found in '{effective_id_column}'. "
 77                "See duplicate_ids.csv.",
 78                duplicate_rows=duplicate_rows,
 79            )
 80        keep = "first" if options.dedupe_id == "first" else "last"
 81        drop_mask = working[effective_id_column].duplicated(keep=keep) & ~blank_id_mask(
 82            working[effective_id_column]
 83        )
 84        excluded = working.loc[drop_mask].copy()
 85        excluded["_exclusion_reason"] = f"Duplicate ID removed by dedupe={keep}"
 86        excluded_parts.append(excluded)
 87        working = working.loc[~drop_mask].copy()
 88        warnings.append(
 89            f"{len(excluded)} duplicate ID row(s) removed using dedupe={keep}."
 90        )
 91
 92    if options.sample_size is not None and options.sample_size <= 0:
 93        raise AuditSamplingError("--sample-size must be a positive integer.")
 94    needs_sample_size = options.method == "random" or (
 95        options.method == "stratified" and bool(options.strata_proportions)
 96    )
 97    if needs_sample_size:
 98        if options.sample_size is None:
 99            raise AuditSamplingError(f"--sample-size is required for {options.method}.")
100        if options.sample_size > len(working) and not (
101            options.method == "stratified" and options.allow_shortfall
102        ):
103            raise AuditSamplingError(
104                "--sample-size cannot exceed the validated population size."
105            )
106
107    strata_counts: dict[str, int] = {}
108    strata_proportions: dict[str, float] = {}
109    if options.method == "stratified":
110        if not options.stratify_column:
111            raise AuditSamplingError("--stratify-column is required for stratified.")
112        has_counts = bool(options.strata_counts)
113        has_proportions = bool(options.strata_proportions)
114        if has_counts == has_proportions:
115            raise AuditSamplingError(
116                "Use exactly one of --strata-counts or --strata-proportions."
117            )
118        strata_counts = parse_key_ints(options.strata_counts, "strata")
119        strata_proportions = parse_key_floats(options.strata_proportions, "strata")
120        requested_strata = set(strata_counts or strata_proportions)
121        actual_strata = set(working[options.stratify_column].dropna().astype(str))
122        missing = sorted(requested_strata - actual_strata)
123        if missing:
124            raise AuditSamplingError(
125                "Requested strata not found in population: " + ", ".join(missing)
126            )
127
128        if strata_counts:
129            _validate_stratum_counts_fit(working, options, strata_counts)
130
131    if excluded_parts:
132        excluded_rows = pd.concat(excluded_parts, ignore_index=True)
133    else:
134        excluded_rows = working.iloc[0:0].copy()
135
136    return ValidationResult(
137        population=working,
138        excluded_rows=excluded_rows,
139        duplicate_rows=duplicate_rows,
140        effective_id_column=effective_id_column,
141        id_column_omitted=id_column_omitted,
142        blank_id_count=blank_id_count,
143        duplicate_id_count=duplicate_id_count,
144        warnings=warnings,
145        strata_counts=strata_counts,
146        strata_proportions=strata_proportions,
147    )
148
149
150def _validate_stratum_counts_fit(population, options, counts: dict[str, int]) -> None:
151    for stratum, requested in counts.items():
152        available = int((population[options.stratify_column] == stratum).sum())
153        if requested > available and not options.allow_shortfall:
154            raise AuditSamplingError(
155                f"Stratum '{stratum}' has {available} rows; requested {requested}. "
156                "Use --allow-shortfall to continue."
157            )