Python: writing to CSV in a for-loop, conditionally adding value in particular column

Question

Here's an example of the contents of my CSV file:

Fruit, colour, ripe,

apple, green,,
banana, yellow,,
pineapple, green,,
plum, purple,,

I want to loop through the contents of the CSV file and according to a test (extrinsic to the CSV data, using an input value supplied to the enclosing function), end up with something like this:

Fruit, colour, ripe,

apple, green, true, 
banana, yellow,, 
pineapple, green,, 
plum, purple, true,

My current code looks like this:

csv_data = csv.reader(open('./data/fruit_data.csv', 'r'))
for row in csv_data:
    fruit = row[0]
    if fruit == input:
    # Here, write 'true' in the 'ripe' column.

It's easy enough to add new data in one go, using the CSV module or pandas, but here I need to add the data iteratively. It seems that I can't change the CSV file in place(?), but if I write out to a different CSV file, it's going to overwrite on each match within the loop, so it'll only reflect that value.

Answer 1

You have, basically, two approaches:

1- Open a second text file before your loop then loop through each row of the initial file and append rows to the second file. After all rows are done, close the initial file. Example: How do you append to a file?

2- Read in everything from the initial csv. Then make changes to the object you created (I highly recommend using Pandas for this). Then write out to a csv. Here's an example of that method:

import pandas as pd
import numpy as np

# read in the csv
csv_data = pd.read_csv('./data/fruit_data.csv')

# I'm partial to the numpy where logic when creating a new column based 
# on if/then logic on an existing column
csv_data['ripe'] = np.where(csv_data['fruit']==input, True, False)

# write out the csv
csv_data.to_csv('./data/outfile.csv')

The choice between 1 and 2 should really come down to scale. If your csv is so big that you can't read it all in and manipulate it the way you want, then you should molest it line by line. If you can read the whole thing in and then manipulate it with Pandas, your life will be MUCH easier.

Answer 2

If you want to create a new CSV file

csv_data = csv.reader(open('./Desktop/fruit_data.csv', 'r'))
csv_new = csv.writer(open('./Desktop/fruit_new_data.csv', 'w'))
for row in csv_data:
    fruit = row[0]
    if fruit == input:
        row.append("ripe")
        csv_new.writerow(row)
    else:
        csv_new.writerow(row)

Basically the only thing missing in your previous question is the last statment, which is to write, else is add in case the criteria does not match.

Another possibility could be to use linestartswith