数据框字典的 Python 单元测试

Question

我有一个返回 pandas 数据框字典的函数，我想为它设计一个单元测试。

我知道如何对 pandas 数据帧进行单元测试:

import pandas as pd
from pandas.util.testing import assert_frame_equal
import unittest

df1 = pd.DataFrame({'a':[1,2,3], 'b': [4,5,6], 'c': [7,8,9]}) 
df2 =pd.DataFrame(df1)

class DictEq(unittest.TestCase):
    def test_dict_eq(self):
        assert_frame_equal(df1, df2)

unittest.main()

但是，我似乎不明白如何设计比较以下内容的测试:

dict1 = {'a': df1}
dict2 = {'a': df2}

我试过以下方法，都失败了:

from nose.tools import assert_equal, assert_dict_equal  

class DictEq(unittest.TestCase):
    def test_dict_eq1(self):
        assert_equal(dict1, dict2)

    def test_dict_eq2(self):
        assert_dict_equal(dict1 , dict2)

    def test_dict_eq3(self):
        self.assertTrue(dict1 == dict2)

pandas.util.testing 的 assert_dict_equal 函数也失败了。

Answer 1

试试这个:

import pandas as pd

df1 = pd.DataFrame({'a':[1,2,3], 'b': [4,5,6], 'c': [7,8,9]})
df2 =pd.DataFrame(df1)

class DfWrap():
    def __init__(self, df):
        self.df = df
    def __eq__(self, df2):
        return all(self.df == df2)


dic1 = {'a': DfWrap(df1)}
dic2 = {'a': DfWrap(df2)}

print(dic1 == dic2)

这输出 True。它也应该与 assert_dict_equal 一起使用，只要您将 dataframe 对象包装在 DfWrap 中。

这就是它起作用的原因:

你必须想象，为了比较字典，python 将遍历每个键(递归)并调用__eq__(或==)进行比较.问题是，当您在 dataframe 上调用 __eq__(或 ==)时，它不会返回 bool。相反，它返回另一个 dataframe:

import pandas as pd

df1 = pd.DataFrame({'a':[1,2,3], 'b': [4,5,6], 'c': [7,8,9]})
df2 =pd.DataFrame(df1)

df_compare = df1 == df2
print(type(df_compare))

这个输出:

<class 'pandas.core.frame.DataFrame'>

因此，包装器使 df1 == df2 输出 bool 而不是 dataframe:

DfWrap(df1) == DfWrap(df2)

评估为 True。

HTH.