如何在Python中拆分csv文件，使其在每个较小的文件中保持头部？

import csv

def split_csv(filename, chunk_size):
    with open(filename, 'r') as file:
        reader = csv.reader(file)
        headers = next(reader)
        current_chunk = 1
        current_rows = []

        for i, row in enumerate(reader, start=1):
            current_rows.append(row)

            if i % chunk_size == 0:
                write_chunk(current_rows, headers, current_chunk)
                current_rows = []
                current_chunk += 1

        # 处理剩余的行
        if current_rows:
            write_chunk(current_rows, headers, current_chunk)

def write_chunk(rows, headers, chunk_number):
    filename = f"chunk_{chunk_number}.csv"

    with open(filename, 'w', newline='') as file:
        writer = csv.writer(file)
        writer.writerow(headers)
        writer.writerows(rows)

# 使用示例
split_csv('input.csv', 1000)