zsjz-ai/docs/ai-clean-e2e/generate_fixtures.py

#!/usr/bin/env python3
"""生成 AI 数据清洗 E2E 验收用的 xlsx/csv 样本。

依赖：openpyxl
用法：python generate_fixtures.py   # 输出到同目录 fixtures/
"""
from __future__ import annotations

import csv
import os

from openpyxl import Workbook
from openpyxl.styles import Font

OUT = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures")
os.makedirs(OUT, exist_ok=True)

HEADERS = ["姓名", "身份证号", "手机号", "银行卡号", "交易日期", "交易时间", "金额", "数量", "备注"]

ROWS = [
    ["张三", "11010519491231002X", "13800000000", "6222020200001234",
     "2026-01-05", "2026-01-05 10:30:00", "1,234.50", "3", "正常"],
    ["李四", "110105491231002", "+86 139-0000-1111", "6222 0202 0000 1235",
     "2026/1/6", "2026-01-06T11:00:00", "-100.00", "2", "负数金额"],
    ["王五", "12345", "123", "6.22202E+15",
     "2026年1月7日", "10:30", "(200.50)", "1.5", "坏证件/电话/卡/纯时间"],
    ["赵六", "110105194912310021", "13800002222", "6222020200001236",
     "2026-02-30", "25:00", "约100元", "abc", "坏日期/时间/金额/数量"],
    ["钱七", "11010519491231002X", "13800003333", "6222020200001237",
     "45000", "1704067200", "¥2,300.5", "10", "序列号/epoch"],
    ["孙八", "11010519491231002X", "13800004444", "6222020200001238",
     "20260108", "2026-01-08 09:00:00", "100元", "8", "紧凑日期/带元"],
    [],
    ["周九", "11010519491231002X", "13800005555", "6222020200001239",
     "2026-01-09", "2026-01-09 12:00:00", "300.00", "5", "空行后的行"],
]


def save_xlsx(name: str, sheet_name: str, rows, headers=None, title=None, start_row=1,
              positions=None):
    wb = Workbook()
    ws = wb.active
    ws.title = sheet_name
    row = start_row
    if title:
        ws.merge_cells(start_row=1, start_column=1, end_row=1, end_column=len(headers or HEADERS))
        ws.cell(row=1, column=1, value=title).font = Font(bold=True)
        row = 3
    for c, h in enumerate(headers or HEADERS, 1):
        ws.cell(row=row, column=c, value=h).font = Font(bold=True)
    for i, values in enumerate(rows):
        target = positions[i] if positions else row + 1 + i
        for c, v in enumerate(values, 1):
            ws.cell(row=target, column=c, value=v)
    wb.save(os.path.join(OUT, name))


def save_csv(name: str, rows):
    with open(os.path.join(OUT, name), "w", newline="", encoding="utf-8") as f:
        w = csv.writer(f)
        w.writerow(HEADERS)
        for values in rows:
            w.writerow(values if values else [""] * len(HEADERS))


def main():
    # 01：自建表类型推断（标题行 + 空行 + 表头 + 空行隔断）
    save_xlsx("01_self_built_types.xlsx", "经营者信息", ROWS, title="经营者信息表")
    # 02：与 01 同表头但倒序（验证复用必须同序）
    save_xlsx("02_reordered_headers.xlsx", "乱序经营者",
              [list(reversed(r)) for r in ROWS[:3]], headers=list(reversed(HEADERS)))
    # 03：两列标题 + 空行 + 表头 + 空行隔断，验证表头定位与 row_no 缺口
    save_xlsx("03_messy_title_blank.xlsx", "人员交易", [
        ["张三", "1,000.00", "2026-01-01"],
        ["李四", "2,000.00", "2026-01-02"],
        ["王五", "3,000.00", "2026-01-03"],
        ["赵六", "坏金额", "坏日期"],
    ], headers=["姓名", "金额", "日期"], title="2026年人员交易明细",
        positions=[4, 5, 7, 8])
    # 04：CSV 版脏数据（验证 CSV 不再 Skip；可选列脏值只置空）
    save_csv("04_csv_types.csv", ROWS)
    # 05：同表头 20 行，1 行身份证非法（配合把 ai_clean_field.required 临时置 1，
    #     验证必填列坏值整行拒绝 + 失败明细）
    rows05 = []
    for i in range(1, 21):
        bad = i == 1
        rows05.append([
            "坏证件行" if bad else f"正常{i}",
            "12345" if bad else "11010519491231002X",
            "13800000000", "6222020200001234", "2026-01-05",
            "2026-01-05 10:30:00", "100.00", "1", "ok",
        ])
    save_xlsx("05_required_bad.xlsx", "经营者信息", rows05)
    # 07：同 05，但坏证件是完整 18 位（校验位错），验证失败明细不脱敏、原值完整展示
    rows07 = []
    for i in range(1, 21):
        bad = i == 1
        rows07.append([
            "坏18位证件行" if bad else f"正常{i}",
            "110105194912310021" if bad else "11010519491231002X",
            "13800000000", "6222020200001234", "2026-01-05",
            "2026-01-05 10:30:00", "100.00", "1", "ok",
        ])
    save_xlsx("07_required_bad_id18.xlsx", "经营者信息", rows07)
    # 06：干净 CSV（验证 CSV 正常入库）
    save_csv("06_csv_clean.csv", [
        [f"CSV正常{i}", "11010519491231002X", "13800000000", "6222020200001234",
         "2026-01-05", "2026-01-05 10:30:00", "100.00", "1", "ok"]
        for i in range(1, 8)
    ])
    print("fixtures written to", OUT)


if __name__ == "__main__":
    main()