Python 3.8.8中read_csv(StringIO())行为为何随字符串定义位置变化？

问题现象 #

将CSV字符串定义在函数外部并传入函数时，pd.read_csv能正常解析生成DataFrame，可通过'code'列名筛选数据；但把CSV字符串移到函数内部定义后，运行会触发KeyError: 'code'。

问题根源 #

函数内部的多行字符串使用了缩进（为了贴合代码格式），导致CSV的每一行开头都带有空格，包括列名行的code前也有多个空格。此时pandas读取后，DataFrame的列名实际是' code'（带前置空格），而非预期的'code'，因此通过df2['code']访问时找不到对应列，抛出KeyError。

而外部定义的CSV字符串所有行左对齐、无前置空格，列名是纯净的'code'，所以能正常访问。

解决方法 #

方法1：移除多行字符串的缩进 #

让CSV内容的每一行都左对齐，不带缩进：

import pandas as pd
from io import StringIO

def test2(search):
    csv_str2 = """\
code,name
AB123,David
CD456,Larry
EF789,Jones
"""
    df2 = pd.read_csv(StringIO(csv_str2))
    result2 = df2.loc[df2['code'].isin(search), ['code']]
    return result2

search_list2 = ['CD456', 'EF789']
print(test2(search_list2))

方法2：使用textwrap.dedent去除公共缩进 #

如果希望保持代码的缩进格式，可借助textwrap模块的dedent函数，自动去除每行开头的公共空白字符：

import pandas as pd
from io import StringIO
import textwrap

def test2(search):
    csv_str2 = textwrap.dedent("""\
    code,name
    AB123,David
    CD456,Larry
    EF789,Jones
    """)
    df2 = pd.read_csv(StringIO(csv_str2))
    result2 = df2.loc[df2['code'].isin(search), ['code']]
    return result2

search_list2 = ['CD456', 'EF789']
print(test2(search_list2))

两种方法都能让pandas正确解析出'code'列，避免KeyError。

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