Python遍历CSV提取指定ID不定数量城市字段值问题求解

问题根源 #

代码报错、逻辑不通过有3个核心问题：

1. 动态键名写法错误：row['city+str(i)'] 会直接查找字面名为city+str(i)的列，不会执行字符串拼接，自然找不到对应字段
2. 列表操作逻辑错误：初始化的列表是cities，循环里误用了未定义的city变量，且append()是列表原地修改方法，不需要给返回值赋值，赋值后会得到None
3. 类型匹配错误：DictReader读取的所有字段值默认是字符串类型，定义的查询ID如果是整数类型，直接相等判断永远不会命中匹配
   另外用len(row)-3计算城市列数的逻辑鲁棒性差，如果后续列顺序调整很容易出错，更稳妥的方式是直接识别所有列名以city开头的字段。

修正后可运行代码 #

from csv import DictReader

target_id = "2"  # 要查询的用户ID，注意和CSV读取的字符串类型匹配
user_name = []
user_age = []
user_cities = []

with open('File_name', 'r') as f:
    # skipinitialspace=True自动处理逗号后的空格，适配示例CSV的表头格式
    reader = DictReader(f, skipinitialspace=True)
    # 提前筛选所有城市列，按city后的数字排序，自动适配后续新增的城市列
    city_columns = sorted(
        [col for col in reader.fieldnames if col.startswith('city')],
        key=lambda x: int(x.replace('city', ''))
    )
    for row in reader:
        if row['ID'] == target_id:
            user_name.append(row['name'])
            user_age.append(row['age'])
            # 遍历所有城市列，过滤空值
            for col in city_columns:
                city_val = row[col]
                if city_val:
                    user_cities.append(city_val)
            # 匹配到目标用户后直接终止遍历，提升效率
            break

print(user_name, user_age, user_cities)

运行效果 #

• 把target_id设为"3"，输出城市列表：["Eureka,NA", "NV,NA"]
• 把target_id设为"2"，输出城市列表：["Los Angeles,NA"]
• 把target_id设为"1"，输出城市列表：["Ann Arbor,NA", "CA,NA", "MS,NA"]

代码会自动识别所有以city开头的列，后续数据源新增city4、city5等列时不需要修改代码，可直接适配。

内容的提问来源于stack exchange，提问作者Stay cool