python写入csv文件中的新列

Question

I want to add a new column to an existing file. But it gets a little complicated with the additional loops i add.

input file:

testfile.csv

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

output i want:

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

This is what i have tried:

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])

And this is the error i got:

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

I was trying to follow this post 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

I couldn't reproduce your error, but i cleaned your code a bit. There is no reason to reopen the input file for every language.

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)