如何使用csv.writer添加自定義列

Question

我有一個python腳本，該腳本需要一個大文件，並且僅復制我想要的幾列，同時還在新文件中寫入新的自定義標頭。

import csv,time,string,os, requests

og_file = "\\\\network\\x\\test.csv"
today = time.strftime("%Y-%m-%d")


#fields to keep from original file
fields = ["As Of Date", "Ph", "Home Country"]

with open(og_file) as infile, open("c:\\upload\\output_" + today + ".csv", "wb") as outfile:
    r = csv.DictReader(infile)
    w = csv.DictWriter(outfile, fields, extrasaction="ignore")

    wtr = csv.writer( outfile )    
    wtr.writerow(["upload_date", "phone", "country"])
    for row in r:
        w.writerow(row)

在編寫此文件時，我需要在每行的開頭創建一個附加列，並從1開始到n一起添加一個數字，因此輸出如下所示：

id |    upload_date    |    phone    |   country
1        2012-01-01         555-1234      USA
2        2012-02-01         555-1235      USA
3        2012-03-01         555-1236      USA
4        2012-04-01         555-1237      USA

我不確定該怎么做。

Answer 1

fields = ["id"] + fields  # so it'll write later

with open(og_file) as infile, open("c:\\upload\\output_" + today + ".csv", "wb") as outfile:
    r = csv.DictReader(infile)
    w = csv.DictWriter(outfile, fields, extrasaction="ignore")

    wtr = csv.writer( outfile )    
    wtr.writerow(["id","upload_date", "phone", "country"])
    for i,row in enumerate(r, start=1):
        row["id"] = i
        w.writerow(row)

Answer 2

您可以為id列使用itertools.count對象：

from itertools import count

with open(og_file) as infile, open("c:\\upload\\output_" + today + ".csv", "wb") as outfile:
    r = csv.DictReader(infile)
    w = csv.DictWriter(outfile, fields, extrasaction="ignore")

    wtr = csv.writer( outfile )    
    wtr.writerow(["id", "upload_date", "phone", "country"])
    c = count(1)
    for row in r:
        row['id'] = next(c)
        w.writerow(row)