Python：如何保留一個文件的標頭並將其附加到另一個文件

Question

我有一個CSV文件data.csv並且已通過以下代碼將此文件分為2個文件test.csv和train.csv ：

import csv
import random

with open('datafile.csv', 'r') as data:
    with open('test.csv', 'w') as test:
        with open('train.csv', 'w') as train:
            test_writer = csv.writer(test)
            train_writer = csv.writer(train)
            for line in csv.reader(data):
                if random.random() > 0.85:
                     test_writer.writerow(line)
                else:
                    train_writer.writerow(line)

這段代碼雖然工作得很好，但頭存在於test.csv ，但沒有發現頭train.csv任何方式更新代碼，並獲得頭在這兩個文件？

Answer 1

最初的問題實際上是可以解決的，但是在進行編輯后，清楚的是要解決的問題：逐行讀取csv.reader並將其隨機寫入任一csv.writer。

您需要像在1.示例中一樣首先閱讀標題。

import csv
import random

with open('datafile.csv', 'r') as data, \
        open('test.csv', 'w') as test, \
        open('train.csv', 'w') as train:
    test_writer = csv.writer(test)
    train_writer = csv.writer(train)
    reader = csv.reader(data)
    header = next(reader)
    test_writer.writerow(header)
    train_writer.writerow(header)
    for row in reader:
        if random.random() > 0.85:
            test_writer.writerow(row)
        else:
            train_writer.writerow(row)