时间戳

首页 > 解决方案 > 数据框字典的 Python 单元测试

python - 数据框字典的 Python 单元测试

问题描述

我有一个返回熊猫数据框字典的函数，我希望为它设计一个单元测试。

我知道如何对 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 函数也失败了。

标签： pythonunit-testing

解决方案

试试这个：

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__（或==）要比较的项目。问题是，当您在 a 上调​​用__eq__（或==）时dataframe，它不会返回 a 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输出 abool而不是 a dataframe：

DfWrap(df1) == DfWrap(df2)

评估为True。

HTH。

推荐阅读