audit-labs/audit-tools

 1"""Input and output helpers for audit sampling."""
 2
 3from __future__ import annotations
 4
 5import hashlib
 6from pathlib import Path
 7
 8import pandas as pd
 9
10SUPPORTED_EXCEL_SUFFIXES = {".xlsx", ".xls", ".xlsm"}
11
12
13class AuditSamplingError(Exception):
14    """Raised when the sampling request cannot be completed."""
15
16    def __init__(self, message: str, **artifacts) -> None:
17        super().__init__(message)
18        self.artifacts = artifacts
19
20
21def sha256_file(path: Path) -> str:
22    digest = hashlib.sha256()
23    with path.open("rb") as handle:
24        for chunk in iter(lambda: handle.read(1024 * 1024), b""):
25            digest.update(chunk)
26    return digest.hexdigest()
27
28
29def load_population(input_path: Path, sheet: str | None = None) -> pd.DataFrame:
30    if not input_path.exists():
31        raise AuditSamplingError(f"Input file does not exist: {input_path}")
32
33    suffix = input_path.suffix.lower()
34    if suffix == ".csv":
35        frame = pd.read_csv(input_path)
36    elif suffix in SUPPORTED_EXCEL_SUFFIXES:
37        excel = pd.ExcelFile(input_path)
38        if sheet is None:
39            if len(excel.sheet_names) != 1:
40                names = ", ".join(excel.sheet_names)
41                raise AuditSamplingError(
42                    "Excel workbook has multiple sheets. Provide --sheet. "
43                    f"Available sheets: {names}"
44                )
45            sheet = excel.sheet_names[0]
46        frame = pd.read_excel(input_path, sheet_name=sheet)
47    else:
48        supported = ".csv, .xlsx, .xls, .xlsm"
49        raise AuditSamplingError(
50            f"Unsupported input extension '{suffix}'. Supported: {supported}"
51        )
52
53    frame = frame.copy()
54    frame.insert(0, "_source_row_number", range(2, len(frame) + 2))
55    return frame
56
57
58def write_csv(frame: pd.DataFrame, path: Path) -> None:
59    frame.to_csv(path, index=False)