将 "assertTrue" 单元测试应用于 pandas DataFrame？

Question

使用 python 的单元测试，如何检查 pandas DataFrame df 的列名是否与列表的列名一致？

import unittest

column_names_list = ["col1","col2","col3"]

assertTrue(df.columns in column_names_list)

例如，我应该为此得到 False df:

col1   col2   col3   col4
aasa   4546   5643   shsh
...

Answer 1

使用 set 你可以找到两个列表之间的区别 assert 那是 0

In [40]: assert len(set(df.columns) - set(['A', 'B', 'C'])) == 0

In [41]: assert len(set(df.columns) - set(['A', 'B', 'D'])) == 0
---------------------------------------------------------------------------
AssertionError                            Traceback (most recent call last)
----> 1 assert len(set(df.columns) - set(['A', 'B', 'D'])) == 0

Answer 2

你也可以使用集合相等。这样断言错误可能更有帮助（对集合使用大括号）：

匹配

df = pd.DataFrame([], columns=['col1', 'col2', 'col3'])
assert set(df.columns) == {'col1', 'col2', 'col3'}

不匹配

df = pd.DataFrame([], columns=['col1', 'col2'])
assert set(df.columns) == {'col1', 'col2', 'col3'}