Function assert_eq

dask/dataframe/utils.py:531–588

(
    a,
    b,
    check_names=True,
    check_dtype=True,
    check_divisions=True,
    check_index=True,
    sort_results=True,
    scheduler="sync",
    **kwargs,
)

def assert_eq(
    a,
    b,
    check_names=True,
    check_dtype=True,
    check_divisions=True,
    check_index=True,
    sort_results=True,
    scheduler="sync",
    **kwargs,
):
    if check_divisions:
        assert_divisions(a, scheduler=scheduler)
        assert_divisions(b, scheduler=scheduler)
        if hasattr(a, "divisions") and hasattr(b, "divisions"):
            at = type(np.asarray(a.divisions).tolist()[0])  # numpy to python
            bt = type(np.asarray(b.divisions).tolist()[0])  # scalar conversion
            assert at == bt, (at, bt)
    assert_sane_keynames(a)
    assert_sane_keynames(b)
    a = _check_dask(
        a, check_names=check_names, check_dtypes=check_dtype, scheduler=scheduler
    )
    b = _check_dask(
        b, check_names=check_names, check_dtypes=check_dtype, scheduler=scheduler
    )
    if hasattr(a, "to_pandas"):
        a = a.to_pandas()
    if hasattr(b, "to_pandas"):
        b = b.to_pandas()

    a, b = _maybe_convert_string(a, b)

    if isinstance(a, (pd.DataFrame, pd.Series)) and sort_results:
        a = _maybe_sort(a, check_index)
        b = _maybe_sort(b, check_index)
    if not check_index:
        a = a.reset_index(drop=True)
        b = b.reset_index(drop=True)
    if isinstance(a, pd.DataFrame):
        tm.assert_frame_equal(
            a, b, check_names=check_names, check_dtype=check_dtype, **kwargs
        )
    elif isinstance(a, pd.Series):
        tm.assert_series_equal(
            a, b, check_names=check_names, check_dtype=check_dtype, **kwargs
        )
    elif isinstance(a, pd.Index):
        tm.assert_index_equal(a, b, exact=check_dtype, **kwargs)
    else:
        if a == b:
            return True
        else:
            if np.isnan(a):
                assert np.isnan(b)
            else:
                assert np.allclose(a, b)
    return True