Python: add value and write output

Question

I need to get information from a list and add a column year from name. I still not sure how to add one field 'year' in record. Can I use append? And about output file, I just need use outputcsv.writerow(records) isn't it? This is a part of code that I stuck:

filenames = ('babyQld2010.csv',     
    'babyQld2011.csv',
    'babyQld2012.csv', 
    'babyQld2012.csv',
    'babyQld2014.csv')

outFile = open('babyQldAll.csv','w') 
csvFile_out = csv.writer(outFile, delimiter=',')

for filename in filenames:
    name, ext = filename.split('.')
    year = name[-4:]     #extract year from  file names
    records = extract_names(filename)

    # Get (name, count, gender)  from  list "records", 
    # and add value of "year" and write into output file (using "for" loop )

Output file look like:

2010,Lola,69,Girl

And input, I have 5 file babyQld2010.csv, babyQld2011.csv, babyQld2012.csv, babyQld2012.csv, babyQld2014.csv which contains:

Mia,425,William,493

and I have to sort it in format and I already done it and save in list 'records'

Lola,69,Girl

now I need to add one field 'year' on 'record' list and export csv file.

This is my full code:

import csv

def extract_names(filename):
    ''' Extract babyname, count, gender from a csv file,
    and return the data in a list.
    '''

    inFile = open(filename, 'rU')
    csvFile = csv.reader(inFile, delimiter=',')

    # Initialization
    records = []
    rowNum = 0

    for row in  csvFile:

        if rowNum != 0:

            # +++++ You code here ++++
            # Read each row of csv file and save information  in list 'records'
            # as (name, count, gender) 
            records.append([row[0], row[1], "Female"])
            records.append([row[2], row[3], "Male"])
            print('Process each row...')      

        rowNum += 1

    inFile.close()
    return(records)



#### Start main program  #####

filenames = ('babyQld2010.csv',     
    'babyQld2011.csv',
    'babyQld2012.csv', 
    'babyQld2012.csv',
    'babyQld2014.csv')

with open('babyQldAll.csv','w') as outFile:

    csvFile_out = csv.writer(outFile, delimiter=',')

    for filename in filenames:
        name, ext = filename.split('.')
        year = name.split('.')[0][-4:]     #extract year from  file names
        records = extract_names(filename)
        for record in records:
            csvFile_out.write([year] + record)
    print("Write in csv file...")     

outFile.close()

Answer 1

Yes, you can just append the year column to each row as you read it in from your source files. You can read in & write out each row as a dictionary so that you can use your existing column headers to address the data if you need to massage it on the way through.

Using the csv.DictWriter() method you specify your headers (fieldnames) when you set it up. You can then write them out with the writeheader() method.

import csv
file_list = ['babyQld2010.csv',     
             'babyQld2011.csv',
             'babyQld2012.csv', 
             'babyQld2012.csv',
             'babyQld2014.csv']
outFile = open('babyQldAll.csv', 'wb')
csv_writer = csv.DictWriter(outFile, 
                            fieldnames=['name','count','gender','year'])
csv_write_out.writeheader()

for a_file in file_list:
    name,ext = a_file.split('.')
    year = name[-4:]
    with open(a_file, 'rb') as inFile:
        csv_read_in = csv.DictReader(inFile)
        for row in csv_read_in:
            row['year'] = year
            csv_writer.writerow(row)

outfile.close()

Hope this helps.

Answer 2

To get the year from the csv file you can simply split the string at '.' and then take the last four characters from the first part of the split. Example -

>>> s = 'babyQld2010.csv'
>>> s.split('.')[0][-4:]
'2010'

Then just simply iterate over your list of records, which you say is correct, for each list within in, use list contatenation to create a new list with year at the start and write that to csv file.

I would also suggest that you use with statement for opening the file to write to (and even in the function where you are reading from the other csv files). Example -

filenames = ('babyQld2010.csv',     
    'babyQld2011.csv',
    'babyQld2012.csv', 
    'babyQld2012.csv',
    'babyQld2014.csv')

with open('babyQldAll.csv','w') as outFile:

    csvFile_out = csv.writer(outFile, delimiter=',')

    for filename in filenames:
        name, ext = filename.split('.')
        year = name.split('.')[0][-4:]     #extract year from  file names
        records = extract_names(filename)
        for record in records:
            csvFile_out.writerow([year] + record)