import numpy as np
import numpy.testing as npt
import pandas as pd
def assert_frames_equal(actual, expected, use_close=False):
"""
Compare DataFrame items by index and column and
raise AssertionError if any item is not equal.
Ordering is unimportant, items are compared only by label.
NaN and infinite values are supported.
Parameters
----------
actual : pandas.DataFrame
expected : pandas.DataFrame
use_close : bool, optional
If True, use numpy.testing.assert_allclose instead of
numpy.testing.assert_equal.
"""
if use_close:
comp = npt.assert_allclose
else:
comp = npt.assert_equal
assert (isinstance(actual, pd.DataFrame) and
isinstance(expected, pd.DataFrame)), \
'Inputs must both be pandas DataFrames.'
for i, exp_row in expected.iterrows():
assert i in actual.index, 'Expected row {!r} not found.'.format(i)
act_row = actual.loc[i]
for j, exp_item in exp_row.iteritems():
assert j in act_row.index, \
'Expected column {!r} not found.'.format(j)
act_item = act_row[j]
try:
comp(act_item, exp_item)
except AssertionError as e:
raise AssertionError(
e.message + '\n\nColumn: {!r}\nRow: {!r}'.format(j, i))
expected = pd.DataFrame({'a': [1, np.nan, 3],
'b': [np.nan, 5, 6]},
index=['x', 'y', 'z'])
actual = pd.DataFrame([[4, 1],
[6, 3],
[5, np.nan]],
index=['x', 'z', 'y'],
columns=['b', 'a'])
assert_frames_equal(actual, actual)
assert_frames_equal(actual, expected)
--------------------------------------------------------------------------- AssertionError Traceback (most recent call last) <ipython-input-57-2fa991ae8dd6> in <module>() ----> 1 assert_frames_equal(actual, expected) <ipython-input-53-fedbc359fc19> in assert_frames_equal(actual, expected) 24 except AssertionError as e: 25 raise AssertionError(e.message + ---> 26 '\n\nColumn: {!r}\nRow: {!r}'.format(j, i)) AssertionError: Items are not equal: ACTUAL: 4.0 DESIRED: nan Column: 'b' Row: 'x'