MCM-ICM_Toolkit/src/mcm_toolkit/data/io.py

"""读写表格。

只保留两个函数：`load_table` / `save_table` 按扩展名分派。
原建议里的 `load_csv` `load_excel` `save_csv` `save_excel` 四个函数
只是在转发 pandas，合并后调用点更短。
"""

from __future__ import annotations

from pathlib import Path
from typing import Any

import pandas as pd

__all__ = ["load_table", "save_table"]


def load_table(path: str | Path, **kwargs: Any) -> Any:
    """按扩展名自动选择读取方式，返回 DataFrame。

    约定：
    * `.csv` 默认做编码探测（utf-8 → utf-8-sig → gbk → gb18030），
      国内赛题数据常见 GBK；
    * `.xlsx` / `.xls` 需要 `openpyxl`（`pip install "mcm-toolkit[excel]"`）；
    * `na_values` 默认包含空串、`"NA"`、`"N/A"`、`"#N/A"`、`"null"`、`"-"`。

    Examples
    --------
    >>> df = load_table("data.csv")          # doctest: +SKIP
    >>> df.shape                             # doctest: +SKIP
    """
    path = Path(path)
    suffix = path.suffix.lower()

    encodings = ['utf-8','utf-8-sig','gbk','gb18030']
    kwargs.setdefault('na_values',['','NA','N/A','#N/A','null','-'])

    if suffix == ".csv":
        if 'encoding' in kwargs:
            return pd.read_csv(path,**kwargs)
        for enc in encodings:
            try:
                return pd.read_csv(path,encoding=enc,**kwargs)
            except UnicodeDecodeError:
                continue
            except FileNotFoundError:
                raise FileNotFoundError(f"在这里找不到文件：{path}") from None
        raise ValueError(f"无法解码文件，目前支持解码方式：{encodings}")
    if suffix in ['.xlsx','.xls']:
        return pd.read_excel(path,**kwargs)
    raise ValueError(f"不支持的文件后缀：{suffix}")


def save_table(df: pd.DataFrame, path: str | Path, index: bool = False, **kwargs: Any) -> Path:
    """保存 DataFrame，返回实际写入的路径。"""
    path = Path(path)
    path.parent.mkdir(parents=True,exist_ok=True)
    suffix = path.suffix.lower()

    if not isinstance(df,pd.DataFrame):
        raise TypeError(f"df 必须是一个pandas DataFrame,但是你上传了{type(df)}")

    if suffix == '.csv':
        df.to_csv(path,index=index,**kwargs)
    elif suffix in ['.xlsx','.xls']:
        df.to_excel(path,index=index,**kwargs)
    else:
        raise ValueError(f"不支持的文件后缀：{suffix}")
    return path