Python3.8读取CSV筛选指定行并生成仅含目标列的新CSV文件问题

完整实现代码 #

import csv

# 定义要筛选的学生姓名集合
TARGET_STUDENTS = {"John", "Lisa"}
# 定义输出文件保留的列
OUTPUT_FIELDS = ["Student Name", "Course Name", "GPA"]

matched_rows = []

# 第一步：读取源文件，筛选匹配记录
with open("My_Source_file.csv", "r", encoding="utf-8") as infile:
    reader = csv.reader(infile)
    # 处理表头，去除每个字段首尾空格
    header = [col.strip() for col in next(reader)]
    # 获取目标列的索引，避免硬编码索引出错
    name_idx = header.index("Student Name")
    course_idx = header.index("Course Name")
    gpa_idx = header.index("GPA")
    
    for row in reader:
        # 处理当前行所有字段的首尾空格
        processed_row = [col.strip() for col in row]
        # 筛选符合姓名要求的记录
        if processed_row[name_idx] in TARGET_STUDENTS:
            # 提取需要的三列加入匹配列表
            matched_rows.append([
                processed_row[name_idx],
                processed_row[course_idx],
                processed_row[gpa_idx]
            ])

# 第二步：判断匹配结果，决定是否生成输出文件
if not matched_rows:
    print("No Student name(s) found in the database")
else:
    # 仅存在匹配记录时才创建输出文件
    with open("My_Output_file.csv", "w", newline="", encoding="utf-8") as outfile:
        writer = csv.writer(outfile)
        # 写入表头
        writer.writerow(OUTPUT_FIELDS)
        # 写入所有匹配行
        writer.writerows(matched_rows)

核心逻辑说明 #

• 修复了原代码索引错误问题：源文件一共5列，索引从0开始最大为4，原代码硬编码row[3],row[4],row[5]会出现索引越界，改为通过表头列名动态获取索引，适配性更强
• 全字段去空格：读取表头和每行数据时都对所有字段执行strip()操作，满足字段去首尾空格的要求
• 筛选逻辑：通过集合存储目标学生姓名，遍历每行时判断处理后的学生姓名是否在集合内，符合条件才保留
• 输出控制：先将所有匹配记录加载到内存中，仅当存在匹配记录时才创建输出文件写入内容，无匹配时直接打印提示，不会生成空文件

内容的提问来源于stack exchange，提问作者Data Engineer