蟒蛇寫入到csv文件

Question

新的專欄中，我想一個新列添加到現有的文件。但是它增加了一些額外的循環。

輸入文件：

testfile.csv

col1,col2,col3 
1,2,3 
3,4,5 
4,6,7 

輸出，我想：

USA_testfile.csv

col1,col2,col3,country 
1,2,3,USA 
3,4,5,USA 
4,6,7,USA 

UK_testfile.csv

col1,col2,col3,country 
1,2,3,UK 
3,4,5,UK 
4,6,7,UK 

這是我曾嘗試：

import csv 
import sys 
country_list= ['USA', 'UK'] 

def add_col(csv_file): 
    for country in country_list: 
     with open(csv_file, 'rb') as fin: 
      with open(country+"_timeline_outfile_"+csv_file, 'wb') as fout: 
       writer = csv.writer(fout, lineterminator='\n') 
       reader = csv.reader(fin) 

       all_rows =[] 
       row = next(reader) 
       row.append('country') 
       all_rows.append(row) 
       print all_rows 

       for row in reader: 
        row.append(country) 
        all_rows.append(row) 
       writer.writerows(all_rows) 

add_col(sys.argv[1]) 

這是我得到的錯誤：

File "write_to_csv.py", line 33, in add_col 
    writer.writerows(all_rows) 
ValueError: I/O operation on closed file 

我試圖按照這個帖子here

Answer 1

import csv 

countries = ['USA', 'UK'] 
data = list(csv.reader(open('testfile.csv', 'rb'))) 

for country in countries: 
    with open('{0}_testfile.csv'.format(country), 'wb') as f: 
     writer = csv.writer(f) 
     for i, row in enumerate(data): 
      if i == 0: 
       row = row + ['country'] 
      else: 
       row = row + [country] 
      writer.writerow(row) 

Answer 2

我無法複製你的錯誤，但我清理了一下你的代碼。 沒有理由重新打開每種語言的輸入文件。

def add_col(csv_file): 
    with open(csv_file, 'rb') as fin: 
     reader = csv.reader(fin) 
     for country in country_list: 
      fin.seek(0) # jump to begin of file again 
      with open(country+"_timeline_outfile_"+csv_file, 'wb') as fout: 
       writer = csv.writer(fout, lineterminator='\n') 
       header = next(reader) 
       header.append('country') 
       writer.writerow(header) 
       for row in reader: 
        row.append(country) 
        writer.writerow(row)