将文件名添加为CSV文件的最后一列

Question

我有一个Python脚本，它会修改一个CSV文件，把文件名加到最后一列：

import sys
import glob

for filename in glob.glob(sys.argv[1]):
    file = open(filename)
    data = [line.rstrip() + "," + filename for line in file]
    file.close()

    file = open(filename, "w")
    file.write("\n".join(data))
    file.close()

不过，它还把文件名加到了文件的表头（第一行）。我想在表头那里加上“ID”这个字符串。有没有人能建议我该怎么做？

Answer 1

你可以试试：

data = [file.readline().rstrip() + ",id"]
data += [line.rstrip() + "," + filename for line in file]

Answer 2

以下是对你当前代码的一些小建议：

• 使用 file 作为变量名并不好，因为这会和内置类型冲突。
• 你可以使用 with 语法来自动关闭文件对象。
• 你是不是想在表头加一个额外的列，比如叫 Filename，而不是在第一行省略一列呢？
• 如果你的文件名中有逗号（或者不太可能的换行符），你需要确保文件名是用引号括起来的——光是直接加上去是没用的。

考虑到最后一点，我建议你使用 csv 模块，这样可以帮你处理引号的问题。例如，你可以试试下面的代码：

import glob
import csv
import sys

for filename in glob.glob(sys.argv[1]):
    data = []
    with open(filename) as finput:
        for i, row in enumerate(csv.reader(finput)):
            to_append = "Filename" if i == 0 else filename
            data.append(row+[to_append])
    with open(filename,'wb') as foutput:
        writer = csv.writer(foutput)
        for row in data:
            writer.writerow(row)

这可能会以稍微不同的方式给数据加引号，所以你可能需要调整一下 csv.reader 和 csv.writer 的引号选项，具体可以参考 csv模块的文档。

另外，你可能有很好的理由选择用 glob 作为参数，而不是直接用命令行中的文件，但这有点让人意外——你需要像这样调用你的脚本 ./whatever.py '*.csv'，而不能直接用 ./whatever.py *.csv。相反，你可以这样做：

for filename in sys.argv[1:]:

... 让 shell 在脚本知道之前先展开你的 glob。

最后一点，你现在的做法有点危险，因为如果在写回同一个文件名时出现任何问题，你会丢失数据。避免这种情况的标准方法是先写入一个临时文件，如果成功了，再把临时文件重命名为原来的文件。所以，你可以把整个过程改写成：

import csv
import sys
import tempfile
import shutil

for filename in sys.argv[1:]:
    tmp = tempfile.NamedTemporaryFile(delete=False)
    with open(filename) as finput:
        with open(tmp.name,'wb') as ftmp:
            writer = csv.writer(ftmp)
            for i, row in enumerate(csv.reader(finput)):
                to_append = "Filename" if i == 0 else filename
                writer.writerow(row+[to_append])
    shutil.move(tmp.name,filename)

Answer 3

你可以看看官方的csv模块。