guzhujushi
b6456d97d8
由用户本人实现（决策 25），Codex 只做审查与补测试。 - 按扩展名分派 .csv / .xlsx / .xls - .csv 默认编码探测 utf-8 -> utf-8-sig -> gbk -> gb18030； 显式传 encoding= 时不再探测 - na_values 默认吃下空串 / NA / N/A / #N/A / null / - - save_table 自动建父目录，返回实际写入的 Path
70 lines
2.4 KiB
Python
70 lines
2.4 KiB
Python
"""读写表格。
|
||
|
||
只保留两个函数：`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
|