Python中for entry in csv_compare循环仅迭代一次的问题排查

问题根源 #

csv.reader()返回的是迭代器对象，迭代器的特性是只能被完整遍历一次。第一次外层循环执行时，内层循环已经把csv_compare的所有条目遍历完毕，后续外层循环再进入内层循环时，csv_compare已经没有剩余数据可以迭代，所以内层循环不再执行。

方案一：将对比数据转为列表（可重复遍历） #

把对比文件的所有条目提前读取到列表中，这样每次外层循环都能从头遍历所有对比条目：

import csv
finalCSV = {}
with open('input.csv', newline='') as csvfile, open('compare.csv', newline='') as keyCSVFile, open('output.csv', 'w' ,newline='') as OutputCSV:
    csv_input = csv.reader(csvfile)
    csv_compare = csv.reader(keyCSVFile)
    # 将对比文件的所有条目转为列表，支持重复遍历
    compare_entries = list(csv_compare)
    csv_output = csv.writer(OutputCSV)
    csv_output.writerow(next(csv_input))

    for row in csv_input:
        # 遍历列表而非原迭代器
        for entry in compare_entries:
            print(row[0] + ' ' + entry[0])
            if row[0] == entry[0]:
                csv_output.writerow(row)
                break
    
print('wait...')

方案二：用集合存储对比Key（更高效） #

如果只需要对比第一列的Key，建议把对比文件的Key存入集合，利用集合O(1)的查找效率替代嵌套循环，大幅提升性能（尤其适合大文件）：

import csv
finalCSV = {}
with open('input.csv', newline='') as csvfile, open('compare.csv', newline='') as keyCSVFile, open('output.csv', 'w' ,newline='') as OutputCSV:
    csv_input = csv.reader(csvfile)
    csv_compare = csv.reader(keyCSVFile)
    # 提取对比文件的所有Key到集合
    compare_keys = {entry[0] for entry in csv_compare}
    csv_output = csv.writer(OutputCSV)
    csv_output.writerow(next(csv_input))

    for row in csv_input:
        if row[0] in compare_keys:
            print(row[0])
            csv_output.writerow(row)
    
print('wait...')

内容的提问来源于stack exchange，提问作者Zac Borders