Python在csv列中搜索字符串并在另一个csv文件中写入行

Question

我对编码还很陌生，但我试图编写一个读取 csv 文件的代码，以检查用户输入的字符串是否在特定列中，然后写出字符串中存在的行单独的 csv 文件。问题是当我运行代码时，它只打印出 1 行而不是其中包含字符串的多行。我也尝试让它不区分大小写，但不知何故它也不起作用。

我现在的代码如下所示：

import csv
shape = input("Type in a shape ")
color = input("Type in a color")
with open("directory.csv", "r") as f:
    csvreader = csv.reader(f, delimiter=",")
    for row in csvreader:
      if shape.lower() in row[7] and color.lower() in row[6] :
         with open("data.csv", "w") as file:
            writer = csv.writer(file, lineterminator = "\n")
            writer.writerow(row)

非常需要任何帮助或建议！

Answer 1

data.csv 每次以'w' 模式打开时都会被截断，因此它只会包含最终匹配项。将 open 语句移到代码的开头可以解决问题。

import csv
shape = input("Type in a shape ")
color = input("Type in a color")
with open("directory.csv", "r") as f, open("data.csv", "w") as file:
    csvreader = csv.reader(f, delimiter=",")
    for row in csvreader:
        if shape.lower() in row[7] and color.lower() in row[6] :
            writer = csv.writer(file, lineterminator = "\n")
            writer.writerow(row)

或者，可以在'a'（追加）模式下打开文件，然后将保留文件中的现有行：

import csv
shape = input("Type in a shape ")
color = input("Type in a color")
with open("directory.csv", "r") as f:
    csvreader = csv.reader(f, delimiter=",")
    for row in csvreader:
      if shape.lower() in row[7] and color.lower() in row[6] :
         with open("data.csv", "a") as file:
            writer = csv.writer(file, lineterminator = "\n")
            writer.writerow(row)

Answer 2

熊猫在这里为您提供帮助。使用下面的片段。

import pandas as pd
df = pd.read_csv("directory.csv")
df1 = df[df["column1"] == shape & df["column2"] == color]
df1.to_csv("found_data.csv")