如何用Pandas检查CSV文件列是否存在且匹配预期数据类型？

问题描述 #

加载CSV到数据库前需要验证列名和对应数据类型，用assert语句检查时明明格式符合预期却抛出错误，代码和报错如下：

import pandas as pd

data = [['tom', 10], ['nick', 15], ['juli', 14]]
df = pd.DataFrame(data, columns=['Name', 'Age']) 
assert df.dtypes.to_dict() == {"Name": str, "Age":int}  # 列格式符合预期仍抛出断言错误

错误信息

---

AssertionError Traceback (most recent call last)
Cell In[28], line 1
----> 1 assert df.dtypes.to_dict() == {"Name": str, "Age":int}

AssertionError:

原因及解决办法 #

问题根源 #

Pandas的df.dtypes返回的是NumPy的dtype对象（比如字符串列对应object类型，整数列对应int64），而不是Python原生的str、int类型，直接用原生类型对比自然不相等。

解决办法1：用NumPy dtype对象对比 #

直接用NumPy对应的dtype来做断言：

import pandas as pd
import numpy as np

data = [['tom', 10], ['nick', 15], ['juli', 14]]
df = pd.DataFrame(data, columns=['Name', 'Age']) 
assert df.dtypes.to_dict() == {"Name": np.dtype('O'), "Age": np.int64}

解决办法2：转换为类型名称字符串对比 #

把Pandas的dtype转换成名称字符串后再对比，更直观：

import pandas as pd

data = [['tom', 10], ['nick', 15], ['juli', 14]]
df = pd.DataFrame(data, columns=['Name', 'Age']) 
# 提取每个列的类型名称
dtype_dict = {col: dtype.name for col, dtype in df.dtypes.items()}
assert dtype_dict == {"Name": "object", "Age": "int64"}

解决办法3：列名+类型双重严格验证 #

如果需要同时保证列名顺序和类型都完全符合预期，可以分开检查：

import pandas as pd
import numpy as np

# 定义预期的列名和对应的dtype
expected_columns = ['Name', 'Age']
expected_dtypes = [np.dtype('O'), np.int64]

data = [['tom', 10], ['nick', 15], ['juli', 14]]
df = pd.DataFrame(data, columns=['Name', 'Age']) 

# 先检查列名是否一致
assert list(df.columns) == expected_columns, "列名与预期不符"
# 再检查每列的数据类型是否匹配
assert all(df.dtypes == expected_dtypes), "数据类型与预期不符"

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