91e6c47fc0
统一 voc_llm 密钥解析与默认模型；向量化改为本地 mlx 模型；更新 README、gitignore 与流水线文档。 Co-authored-by: Cursor <cursoragent@cursor.com>
103 lines
3 KiB
Python
103 lines
3 KiB
Python
"""
|
||
合并指定目录下所有表头一致的 CSV 文件为一个 CSV。
|
||
"""
|
||
|
||
from __future__ import annotations
|
||
|
||
import argparse
|
||
from pathlib import Path
|
||
|
||
import pandas as pd
|
||
|
||
_PROJECT_ROOT = Path(__file__).resolve().parent
|
||
DEFAULT_INPUT_DIR = _PROJECT_ROOT / '/Users/onesvmwhoops/Documents/no scratch spray for cats'
|
||
DEFAULT_OUTPUT_PATH = _PROJECT_ROOT / "merged_reviews.csv"
|
||
|
||
|
||
def merge_csv_directory(
|
||
input_dir: str | Path,
|
||
output_path: str | Path,
|
||
*,
|
||
pattern: str = "*.csv",
|
||
recursive: bool = False,
|
||
) -> pd.DataFrame:
|
||
"""
|
||
读取目录内所有匹配的 CSV（表头须一致），纵向合并后写入 output_path。
|
||
|
||
:param input_dir: 含多个 CSV 的目录
|
||
:param output_path: 合并结果输出路径
|
||
:param pattern: 文件名 glob，默认 *.csv
|
||
:param recursive: 是否包含子目录中的 CSV
|
||
:return: 合并后的 DataFrame
|
||
"""
|
||
input_dir = Path(input_dir).resolve()
|
||
output_path = Path(output_path).resolve()
|
||
|
||
if not input_dir.is_dir():
|
||
raise NotADirectoryError(f"目录不存在: {input_dir}")
|
||
|
||
globber = input_dir.rglob if recursive else input_dir.glob
|
||
csv_files = sorted(
|
||
p for p in globber(pattern) if p.is_file() and p.resolve() != output_path
|
||
)
|
||
|
||
if not csv_files:
|
||
raise FileNotFoundError(f"未在 {input_dir} 找到匹配 {pattern!r} 的 CSV 文件")
|
||
|
||
frames: list[pd.DataFrame] = []
|
||
expected_columns: list[str] | None = None
|
||
|
||
for path in csv_files:
|
||
df = pd.read_csv(path, dtype=str, keep_default_na=False)
|
||
if expected_columns is None:
|
||
expected_columns = list(df.columns)
|
||
elif list(df.columns) != expected_columns:
|
||
raise ValueError(
|
||
f"表头不一致: {path.name}\n"
|
||
f" 期望: {expected_columns}\n"
|
||
f" 实际: {list(df.columns)}"
|
||
)
|
||
frames.append(df)
|
||
|
||
merged = pd.concat(frames, ignore_index=True)
|
||
output_path.parent.mkdir(parents=True, exist_ok=True)
|
||
merged.to_csv(output_path, index=False, encoding="utf-8-sig")
|
||
|
||
print(f"已合并 {len(csv_files)} 个文件，共 {len(merged)} 行 -> {output_path}")
|
||
for path in csv_files:
|
||
print(f" - {path.name}")
|
||
|
||
return merged
|
||
|
||
|
||
def main() -> None:
|
||
parser = argparse.ArgumentParser(description="合并目录内表头相同的 CSV 文件")
|
||
parser.add_argument(
|
||
"-i",
|
||
"--input",
|
||
type=Path,
|
||
default=DEFAULT_INPUT_DIR,
|
||
metavar="DIR",
|
||
help=f"输入目录（默认: {DEFAULT_INPUT_DIR}）",
|
||
)
|
||
parser.add_argument(
|
||
"-o",
|
||
"--output",
|
||
type=Path,
|
||
default=DEFAULT_OUTPUT_PATH,
|
||
metavar="FILE",
|
||
help=f"输出 CSV 路径（默认: {DEFAULT_OUTPUT_PATH}）",
|
||
)
|
||
parser.add_argument(
|
||
"-r",
|
||
"--recursive",
|
||
action="store_true",
|
||
help="是否递归搜索子目录",
|
||
)
|
||
args = parser.parse_args()
|
||
|
||
merge_csv_directory(args.input, args.output, recursive=args.recursive)
|
||
|
||
|
||
if __name__ == "__main__":
|
||
main()
|