根据文件名添加CSV列

Question

我在一个目录中有以下CSV文件作为输入文件：

USA_Chicago_data.csv
USA_Dallas_data.csv

其数据如下：

Area          Population   No_of_houses
Brookfield    100000        20000
Springdale    50000         10000

我想在文件的开头添加一列，根据文件名使用城市名称。

USA_Chicago_data.csv:的样本输出

City     Area          Population   No_of_houses
Chicago  Brookfield    100000        20000
Chicago  Springdale    50000         10000

它应该有条件读取文件名，从中提取城市并将其添加到第一列中。

我正在尝试以下脚本：

with open(file, 'r') as csvinput:
if ("Chicago" in file):
with open('file'+_output.csv', 'w') as csvoutput:
writer = csv.writer(csvoutput, lineterminator='n')
for row in csv.reader(csvinput):
writer.writerow(row + ['City'])

但它正在陷入错误。解决方案是什么？我想在不使用熊猫的情况下实现这一目标。

Answer 1

这样的东西应该能起到作用。您在第二个"；打开"；语句，因为引用混淆。

# Extract the city:
city = file.split('_')[1]
with open(file,'r') as csvinput, open(f'USA_{city}_data_output.csv','w') as cvsoutput:
writer = csv.writer(csvoutput, lineterminator='n')
for row in csv.reader(csvinput):
if row[0] == 'Area':
row.insert( 0, 'City')
else:
row.insert( 0, city )
writer.writerow(row)

跟进

要覆盖原始文件，您需要：

# Extract the city:
city = file.split('_')[1]
newfile = f'USA_{city}_data_output.csv'
with open(file,'r') as csvinput, open(newfile,'w') as cvsoutput:
writer = csv.writer(csvoutput, lineterminator='n')
for row in csv.reader(csvinput):
if row[0] == 'Area':
row.insert( 0, 'City')
else:
row.insert( 0, city )
writer.writerow(row)
os.remove( file )
os.rename( newfile, file )