python - 使用 csv.reader 和 csv.writer 写入同一文件而不是不同的文件？

Question

如何将以下代码转换为写入同一文件，而不必写入新文件？

import csv
import itertools
from operator import itemgetter

def filter_data(data):
    for classname, group in itertools.groupby(data, itemgetter(2)):
        filtered_group = [line for line in group]
        new_count = len(filtered_group)
        for line in filtered_group:
            if line[5] == "Count":
                line[5] = "Counter"
            else:
                line[5] = new_count
            yield line

with open('main.csv', 'rb') as f_in, open('main1.csv', 'wb') as f_out:
    reader = csv.reader(f_in)
    writer = csv.writer(f_out)
    writer.writerows(filter_data(reader))

如果 stringio 不起作用，我的尝试...

import csv
import itertools
from operator import itemgetter
import StringIO

def filter_data(data):
    for classname, group in itertools.groupby(data, itemgetter(2)):
        filtered_group = [line for line in group]
        new_count = len(filtered_group)
        for line in filtered_group:
            if line[5] == "Count":
                line[5] = "Counter"
            else:
                line[5] = new_count
            yield line

output = StringIO.StringIO()

with open('../main.csv', 'rb') as f_in:
    reader = csv.reader(f_in)
    output.writelines(filter_data(reader))

contents = output.getvalue()
output.close()

with open('../main_test.csv', 'wb') as f_out: 
    f_out.writelines(contents)

Answer 1

尝试以下操作:

代码

import csv
import itertools
from operator import itemgetter
import StringIO

def filter_data(data):
    for classname, group in itertools.groupby(data, itemgetter(2)):
        filtered_group = [line for line in group]
        new_count = len(filtered_group)
        for line in filtered_group:
            if line[5] == "Count":
                line[5] = "Counter"
            else:
                line[5] = new_count
            yield ','.join(map(str, line)) + '\n'

output = StringIO.StringIO()

with open('../main.csv', 'rb') as f_in:
    reader = csv.reader(f_in)
    output.writelines(filter_data(reader))

contents = output.getvalue()
output.close()

with open('../main.csv', 'wb') as f_out: 
    f_out.writelines(contents)

这应该有效:)