Python 利用CSV模块处理数据

CSV（Comma-Separated Values 逗号分割值）是一种简单的数据存储与分享方式，和Excel相比，CSV文件的一个主要优点是有很多程序可以存储，转换和处理纯文本文件，因此应用场景更加广泛。

一、CSV模块简介

考虑下面的CSV文件：salary.csv。虽然可以使用str.split(‘,’)这样的方式处理数据，但是当遇到带逗号分隔的金额会导致拆分错误，除此之外还可能遇到其他各种问题：

id,name,hire_date, salary ,bonus
1,Vincent,2024/1/1," 1,000.00 ",100
2,Victor,2024/1/2," 2,000.00 ",200
3,Grace,2024/1/3," 3,000.00 ",300

使用split拆分代码如下，金额1,000.00被错误的分成了2段：

input_file = 'salary.csv'
with open(input_file, 'r', newline='') as fi:
        header = fi.readline()
        header = header.strip()
        header_list = header.split(',')
        print(header_list)
        for row  in fi:
            field_list = row.strip().split(',')
            print(field_list)

对于这类问题，建议选择标准库中的CSV模块来进行处理，CSV模块被设计成能够识别微软Excel所采用的编码规则，可以识别出这种模式并正确的分析数据。

CSV模块常用函数如下：

• csv.reader(csvfile, dialect=‘excel’, **fmtparams)，返回一个reader对象，通过该对象处理csvfile中的数据，此函数用来打开输入文件。csvfile必须是包含字符串的可迭代对象，通常是文件型对象或列表，如果csvfile是一个文件对象，则打开时应设置newline=‘’。
• csv.writer(csvfile, dialect=‘excel’, **fmtparams)，返回一个writer对象，该对象负责将数据在给定的文件型对象上转换为带分隔符的字符串，此函数用来打开输出文件。如果csvfile是一个文件对象，则打开时应设置newline=‘’。

二、用法示例

下面两个示例演示从CSV文件中选取特定的行、列，更复杂的操作都可以在此基础上进行扩展。

2.1 使用CSV模块选取特定的行

对于文件salary.csv，过滤出salary列大于1000的行，并写入results.csv。由于原数据类型为字符串且包含逗号，所以需要先去除值里的逗号，最后用float函数转换为浮点数与1000进行比较。筛选出符合条件的数据后调用writer对象上的writerow方法逐行写入results.csv：

import csv
input_file = 'salary.csv'
output_file = 'results.csv'
with open(input_file, 'r', newline='') as fi:
    with open(output_file, 'w', newline='') as fo:
        reader = csv.reader(fi, delimiter=',')    # 默认分隔符就是逗号，delimiter=','也可省略
        writer = csv.writer(fo)
        header = next(reader)    # 读取第一行标题
        print(header)
        writer.writerow(header)    # 先将标题写入文件
        for row in reader:
            salary = float(row[3].replace(',', ''))    # 去除逗号并转换为float型
            if salary > 1000:
                print(row)
                writer.writerow(row)

这里可以看到CSV模块可以正确处理金额中的逗号，不会拆分为2段。

2.2 优化：通过标题名引用列

上面的示例中，在判断金额大于1000时，因为要特殊处理，所以我们用row[3]这种索引方式来选取salary列。如果要处理的列很多，那么都通过row[idx]来引用不仅麻烦而且容易搞混淆。这里用命名元组（collections.namedtuple）优化一下，命名元组赋予每个位置一个含义，提供可读性并添加了通过列名获取值的能力（用索引依然也是可以的）。

import csv
from collections import namedtuple
input_file = 'salary.csv'
output_file = 'results.csv'
with open(input_file, 'r', newline='') as fi:
    with open(output_file, 'w', newline='') as fo:
        reader = csv.reader(fi, delimiter=',')    # 默认分隔符就是逗号，delimeter=','也可省略
        writer = csv.writer(fo)
        header = next(reader)    # 读取第一行标题
        print(header)
        writer.writerow(header)    # 先将标题写入文件
        
        namedrow = namedtuple('namedrow', header)    #  命名元组定义
        for row in reader:
            row = namedrow(*row)    # 转换为命名元组
            salary = float(row.salary.replace(',', ''))    # 通过row.salary引用salary列
            if salary > 1000:
                print(row)
                writer.writerow(row)

2.3 使用CSV模块筛选特定的行

现只需要选取id, name, salary三列。处理思路为通过标题先将3列的索引位置保存到列表中，然后在每行循环中，根据索引位置只保存这3列数据：

import csv
input_file = 'salary.csv'
output_file = 'results.csv'
with open(input_file, 'r', newline='') as fi:
    with open(output_file, 'w', newline='') as fo:
        reader = csv.reader(fi, delimiter=',')    # 默认分隔符就是逗号，delimiter=','也可省略
        writer = csv.writer(fo)
        header = next(reader)    # 读取第一行标题

        columns = ['id', 'name', 'salary']    # 需要的列
        col_idx = []
        for idx in range(len(header)):
            if header[idx] in columns:    # 如果在选取的范围内，保存索引
                col_idx.append(idx)               
        writer.writerow(columns)    # 先将标题写入文件

        for row in reader:
            results = []
            for idx in col_idx:
                results.append(row[idx])    # 通过索引只保留需要的列
            print(results)
            writer.writerow(results)