scan_hf_dai_columns.py

#!/usr/bin/env python3
from __future__ import annotations

import json
import math
import os
import sys
from collections import OrderedDict
from pathlib import Path
from typing import Any

import joblib
import pandas as pd
from huggingface_hub import HfApi, hf_hub_download


REPO_ID = "gui-sparim/repositorio_mesa"
OUTPUT_XLSX = "repositorio_mesa_colunas_dai.xlsx"
SUMMARY_JSON = "repositorio_mesa_colunas_dai_summary.json"


def normalize_value(value: Any) -> str | None:
    if value is None:
        return None
    if isinstance(value, float) and math.isnan(value):
        return None
    try:
        if pd.isna(value):
            return None
    except Exception:  # noqa: BLE001
        pass
    if isinstance(value, bytes):
        value = value.decode("utf-8", errors="replace")
    if isinstance(value, (dict, list, tuple, set)):
        value = json.dumps(value, ensure_ascii=False)
    text = str(value).strip()
    return text or None


def first_nonblank_samples(frame: pd.DataFrame, column: str, limit: int = 5) -> list[str]:
    samples: list[str] = []
    for value in frame[column].tolist():
        normalized = normalize_value(value)
        if normalized is not None:
            samples.append(normalized)
            if len(samples) >= limit:
                break
    return samples


def walk_dataframes(node: Any, path: str = "root", seen: set[int] | None = None) -> list[tuple[str, pd.DataFrame]]:
    if seen is None:
        seen = set()

    node_id = id(node)
    if node_id in seen:
        return []
    seen.add(node_id)

    results: list[tuple[str, pd.DataFrame]] = []
    if isinstance(node, pd.DataFrame):
        return [(path, node)]

    if isinstance(node, dict):
        for key, value in node.items():
            results.extend(walk_dataframes(value, f"{path}.{key}", seen))
        return results

    if isinstance(node, (list, tuple)):
        for index, value in enumerate(node):
            results.extend(walk_dataframes(value, f"{path}[{index}]", seen))
        return results

    attrs = getattr(node, "__dict__", None)
    if isinstance(attrs, dict):
        for key, value in attrs.items():
            results.extend(walk_dataframes(value, f"{path}.{key}", seen))

    return results


def main() -> None:
    token = os.environ.get("HF_TOKEN") or os.environ.get("HUGGINGFACE_HUB_TOKEN")
    if not token:
        raise SystemExit("HF_TOKEN or HUGGINGFACE_HUB_TOKEN is required.")

    os.environ.setdefault("HF_HUB_DISABLE_PROGRESS_BARS", "1")

    api = HfApi(token=token)
    tree = list(api.list_repo_tree(REPO_ID, repo_type="dataset", recursive=True, expand=False))
    dai_files = sorted(
        item.path
        for item in tree
        if getattr(item, "type", "file") == "file" and item.path.lower().endswith(".dai")
    )

    ordered_columns: OrderedDict[str, list[str]] = OrderedDict()
    sources: OrderedDict[str, dict[str, str]] = OrderedDict()
    dataframe_rows: list[dict[str, Any]] = []
    scanned_files: list[str] = []
    skipped_files: list[dict[str, str]] = []

    for dai_path in dai_files:
        print(f"Scanning {dai_path}...", file=sys.stderr)
        try:
            local_path = hf_hub_download(
                repo_id=REPO_ID,
                repo_type="dataset",
                filename=dai_path,
                token=token,
            )
            obj = joblib.load(local_path)
            dataframes = walk_dataframes(obj)
            scanned_files.append(dai_path)

            for dataframe_path, frame in dataframes:
                dataframe_rows.append(
                    {
                        "dai_file": dai_path,
                        "dataframe_path": dataframe_path,
                        "row_count": int(frame.shape[0]),
                        "column_count": int(frame.shape[1]),
                    }
                )
                for column in frame.columns:
                    if column in ordered_columns:
                        continue
                    samples = first_nonblank_samples(frame, column)
                    ordered_columns[column] = samples
                    sources[column] = {
                        "dai_file": dai_path,
                        "dataframe_path": dataframe_path,
                    }
        except Exception as exc:  # noqa: BLE001
            skipped_files.append({"dai_file": dai_path, "error": str(exc)})
            print(f"Skipped {dai_path}: {exc}", file=sys.stderr)

    if not ordered_columns:
        raise SystemExit("No dataframe columns were extracted from the .dai files.")

    workbook_data = {
        column: ordered_columns[column] + [""] * (5 - len(ordered_columns[column]))
        for column in ordered_columns
    }
    columns_df = pd.DataFrame(workbook_data)
    origins_df = pd.DataFrame(
        [
            {
                "column_name": column,
                "first_seen_dai_file": source["dai_file"],
                "first_seen_dataframe_path": source["dataframe_path"],
                "sample_count": len(ordered_columns[column]),
            }
            for column, source in sources.items()
        ]
    )
    dataframes_df = pd.DataFrame(dataframe_rows)
    skipped_df = pd.DataFrame(skipped_files or [{"dai_file": "", "error": ""}])

    output_path = Path(OUTPUT_XLSX).resolve()
    summary_path = Path(SUMMARY_JSON).resolve()

    with pd.ExcelWriter(output_path, engine="openpyxl") as writer:
        columns_df.to_excel(writer, sheet_name="colunas", index=False)
        origins_df.to_excel(writer, sheet_name="origem_colunas", index=False)
        dataframes_df.to_excel(writer, sheet_name="dataframes_lidos", index=False)
        skipped_df.to_excel(writer, sheet_name="arquivos_pulados", index=False)

    summary = {
        "repo_id": REPO_ID,
        "output_xlsx": str(output_path),
        "summary_json": str(summary_path),
        "dai_files_found": len(dai_files),
        "dai_files_scanned": len(scanned_files),
        "dai_files_skipped": len(skipped_files),
        "dataframes_found": len(dataframe_rows),
        "unique_columns": len(ordered_columns),
        "columns": list(ordered_columns.keys()),
        "sources": sources,
        "skipped_files": skipped_files,
    }
    summary_path.write_text(json.dumps(summary, ensure_ascii=False, indent=2), encoding="utf-8")
    print(json.dumps(summary, ensure_ascii=False, indent=2))


if __name__ == "__main__":
    main()