query-risk-user-features/scripts/prepare_input.py

#!/usr/bin/env python3
"""Normalize arbitrary MID and anchor-date inputs for the risk feature pipeline."""

from __future__ import annotations

import argparse
import sys
from pathlib import Path

import pandas as pd


PROJECT_DIR = Path(__file__).resolve().parents[1]
if str(PROJECT_DIR) not in sys.path:
    sys.path.insert(0, str(PROJECT_DIR))

from src.user_input import load_users, normalize_users  # noqa: E402


def parse_user(value: str) -> dict[str, str]:
    parts = value.split(":", 3)
    if len(parts) < 2:
        raise ValueError("--user格式应为 MID:yyyyMMdd[:用户类型[:子分类]]")
    return {
        "mid": parts[0].strip(),
        "anchor_dt": parts[1].strip(),
        "user_type": parts[2].strip() if len(parts) > 2 else "待分类",
        "sub_category": parts[3].strip() if len(parts) > 3 else "待分类",
    }


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--input", type=Path, help="已有CSV，至少包含MID和锚点日期")
    parser.add_argument(
        "--user",
        action="append",
        default=[],
        help="可重复传入：MID:yyyyMMdd[:用户类型[:子分类]]",
    )
    parser.add_argument("--lookback-days", type=int, default=180)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()

    if bool(args.input) == bool(args.user):
        parser.error("--input和--user必须且只能使用一种")
    users = (
        load_users(args.input, args.lookback_days)
        if args.input
        else normalize_users(
            pd.DataFrame([parse_user(value) for value in args.user]),
            lookback_days=args.lookback_days,
        )
    )
    args.output.parent.mkdir(parents=True, exist_ok=True)
    users.to_csv(args.output, index=False)
    print(f"users={len(users)} unique_mids={users['mid'].nunique()} output={args.output}")


if __name__ == "__main__":
    main()