CSV文件写入,需要将特定行写入新的csv文件

时间:2017-04-13 01:38:55

标签: python csv

我必须编写一个代码,该代码采用csv文件并从泰坦尼克号中提取与乘客数据相关的数据。我需要从这个文件中取出并写一个新文件,其中包含幸存的第三类乘客(只有这个)和标题。

到目前为止,我已经提供了我的代码(文本中)。它适用于测试用例(它打印#5),但我相信我的target_data_file是空的吗?

我正在研究如何将这些特定行写入target_data_file。我认为它应该是一个for循环的东西 如果survived == str(1) and pclass == str(3),请写信至Target_data_file

虽然不确定!

谢谢!

import csv
from copy import deepcopy

def third_survived(source_data_file, target_data_file):
    """(str, str) -> int
    Input: Source data is the name of a .csv file containing a subset of the 
    Titanic passenger data, and target_data, the name of a new csv file to be 
    created.
    Output: This function will create a new .csv file named target_data_file and 
    write in it the lines from source_data_file that correspond to the third class 
    passengers who survived the sinking. The function returns the number of lines 
    written to target_data_file.

    >>>third_survived('titanic_some.csv', 'titanic_target.csv')
    5
    """

    with open (str(source_data_file), 'r') as file:
        data_reader=csv.reader(file)
        data_orig=[]
        for row in data_reader:
            data_orig.append(row)

    count= 0
    for elements in range(1,len(data_orig)):
        survived=data_orig[elements][1]
        pclass=data_orig[elements][2]
        if survived == str(1) and pclass == str(3):
            count +=1

    with open(str(target_data_file), 'w') as newfile:
        data_writer=csv.writer(newfile)


        if count == 0:
            return data_orig[0]
        else:
            return count

1 个答案:

答案 0 :(得分:0)

您可以随计数循环一起写入target_data_file(并且您不需要data_orig列表)。 那就是:

def third_survived(source_data_file, target_data_file):
    count= 0
    with open (str(source_data_file), 'r') as file:
        data_reader=csv.reader(file)
        with open(str(target_data_file), 'w') as newfile:
            data_writer=csv.writer(newfile)
            for row in data_reader:
                survived=row [1]
                pclass=row [2]
                if survived == "1" and pclass == "3":
                    count +=1
                    data_writer.writerow(row)

    return count

如果您仍然热衷于在count为零时返回第一行(与您的文档相矛盾) - 您可以添加

first_row = None

count的定义之前,并在每次迭代检查

if first_row is None:
    first_row = row

最后回归

if count == 0:
    return first_row
return count