Python yield函数执行两次：CSV分片后最后文件重复生成问题

问题 #

需求是读取CSV文件，将其分片为多个包含完整记录的CSV文件（不可拆分记录），但运行代码后发现最后一个文件被重复创建两次，且两次写入的内容长度完全一致。

原代码 #

def read_in_chunks(file_path, chunk_size):

    with open(file_path, 'rb') as f:

        pending = b''

        while True:

            chunk = f.read(chunk_size)
            if not chunk:
                print('EOF')
                if pending != b'':
                    yield pending
                pending = b''
                break

            print("CHUNK SIZE: " + str(len(chunk)))
            data = pending + chunk
            x1 = data.rfind(b'\n')
            if x1 != -1:
                x2 = len(data) - x1 - 1
                pending = data[-x2:]
                yield data[:x1 + 1]

table_name = 'my_table'
chunk_num = 1
for records_chunk in read_in_chunks('input_file.csv', chunk_size):
    chunk_filename = f'{table_name}_{chunk_num}.csv'
    with open(chunk_filename, 'wb') as chunk_file:
        print("CHUNK LENGTH TO WRITE:" + str(len(records_chunk)))
        chunk_file.write(records_chunk)
    
    chunk_num += 1

运行输出 #

CHUNK SIZE: 10000
CHUNK LENGTH TO WRITE:9998 
CHUNK SIZE: 10000
CHUNK LENGTH TO WRITE:9984 
CHUNK SIZE: 10000
CHUNK LENGTH TO WRITE:9982 
CHUNK SIZE: 10000
CHUNK LENGTH TO WRITE:10022
CHUNK SIZE: 10000
CHUNK LENGTH TO WRITE:9982 
CHUNK SIZE: 3559
CHUNK LENGTH TO WRITE:3591 
EOF
CHUNK LENGTH TO WRITE:3591 

问题原因 #

问题出在pending变量的赋值逻辑上：
当处理最后一个非空chunk时，如果data的最后一个字节恰好是换行符，此时x1 = data.rfind(b'\n')会返回len(data)-1（即最后一个字符的索引）。
随后计算x2 = len(data) - x1 -1，结果为0。而Python中data[-0:]等价于data[0:]，也就是整个data内容，这就导致pending被错误地赋值为当前已经yield过的完整data，而不是空字节串。
当循环进入EOF分支时，pending不为空，会再次yield相同的内容，最终导致最后一个文件被重复写入。

修复方案 #

直接通过切片索引赋值pending，替代原有的x2计算逻辑，这样当x1+1超出data长度时，切片结果自然为空字节串：

修正后的代码 #

def read_in_chunks(file_path, chunk_size):

    with open(file_path, 'rb') as f:

        pending = b''

        while True:

            chunk = f.read(chunk_size)
            if not chunk:
                print('EOF')
                if pending != b'':
                    yield pending
                pending = b''
                break

            print("CHUNK SIZE: " + str(len(chunk)))
            data = pending + chunk
            x1 = data.rfind(b'\n')
            if x1 != -1:
                # 直接取换行符之后的部分作为pending，无需计算x2
                pending = data[x1+1:]
                yield data[:x1 + 1]

table_name = 'my_table'
chunk_num = 1
for records_chunk in read_in_chunks('input_file.csv', chunk_size):
    chunk_filename = f'{table_name}_{chunk_num}.csv'
    with open(chunk_filename, 'wb') as chunk_file:
        print("CHUNK LENGTH TO WRITE:" + str(len(records_chunk)))
        chunk_file.write(records_chunk)
    
    chunk_num += 1

这个修改能确保当最后一个字符是换行符时，pending被设置为空，EOF分支不会重复yield内容，也就不会创建重复的最后一个文件。

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