diff --git a/src/mcm_toolkit/data/io.py b/src/mcm_toolkit/data/io.py index c381f62..0c7568d 100644 --- a/src/mcm_toolkit/data/io.py +++ b/src/mcm_toolkit/data/io.py @@ -10,6 +10,8 @@ from __future__ import annotations from pathlib import Path from typing import Any +import pandas as pd + __all__ = ["load_table", "save_table"] @@ -27,11 +29,41 @@ def load_table(path: str | Path, **kwargs: Any) -> Any: >>> df = load_table("data.csv") # doctest: +SKIP >>> df.shape # doctest: +SKIP """ - # TODO(L0): 编码探测 + 分派;Excel 缺 openpyxl 时给出可读报错。 - raise NotImplementedError + 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: Any, path: str | Path, index: bool = False, **kwargs: Any) -> Path: +def save_table(df: pd.DataFrame, path: str | Path, index: bool = False, **kwargs: Any) -> Path: """保存 DataFrame,返回实际写入的路径。""" - # TODO(L0) - raise NotImplementedError + 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