Python pandas.util.testing.assert_extension_array_equal() Examples

The following are 30 code examples of pandas.util.testing.assert_extension_array_equal(). You can vote up the ones you like or vote down the ones you don't like, and go to the original project or source file by following the links above each example. You may also want to check out all available functions/classes of the module pandas.util.testing , or try the search function .
Example #1
Source File: test_datetimes.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 6 votes vote down vote up 
def test_fillna_preserves_tz(self, method):
        dti = pd.date_range('2000-01-01', periods=5, freq='D', tz='US/Central')
        arr = DatetimeArray(dti, copy=True)
        arr[2] = pd.NaT

        fill_val = dti[1] if method == 'pad' else dti[3]
        expected = DatetimeArray._from_sequence(
            [dti[0], dti[1], fill_val, dti[3], dti[4]],
            freq=None, tz='US/Central'
        )

        result = arr.fillna(method=method)
        tm.assert_extension_array_equal(result, expected)

        # assert that arr and dti were not modified in-place
        assert arr[2] is pd.NaT
        assert dti[2] == pd.Timestamp('2000-01-03', tz='US/Central')
Example #2
Source File: test_integer.py    From recruit with Apache License 2.0 6 votes vote down vote up 
def test_integer_array_constructor():
    values = np.array([1, 2, 3, 4], dtype='int64')
    mask = np.array([False, False, False, True], dtype='bool')

    result = IntegerArray(values, mask)
    expected = integer_array([1, 2, 3, np.nan], dtype='int64')
    tm.assert_extension_array_equal(result, expected)

    with pytest.raises(TypeError):
        IntegerArray(values.tolist(), mask)

    with pytest.raises(TypeError):
        IntegerArray(values, mask.tolist())

    with pytest.raises(TypeError):
        IntegerArray(values.astype(float), mask)

    with pytest.raises(TypeError):
        IntegerArray(values)
Example #3
Source File: test_datetimes.py    From coffeegrindsize with MIT License 6 votes vote down vote up 
def test_fillna_preserves_tz(self, method):
        dti = pd.date_range('2000-01-01', periods=5, freq='D', tz='US/Central')
        arr = DatetimeArray(dti, copy=True)
        arr[2] = pd.NaT

        fill_val = dti[1] if method == 'pad' else dti[3]
        expected = DatetimeArray._from_sequence(
            [dti[0], dti[1], fill_val, dti[3], dti[4]],
            freq=None, tz='US/Central'
        )

        result = arr.fillna(method=method)
        tm.assert_extension_array_equal(result, expected)

        # assert that arr and dti were not modified in-place
        assert arr[2] is pd.NaT
        assert dti[2] == pd.Timestamp('2000-01-03', tz='US/Central')
Example #4
Source File: test_integer.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 6 votes vote down vote up 
def test_integer_array_constructor():
    values = np.array([1, 2, 3, 4], dtype='int64')
    mask = np.array([False, False, False, True], dtype='bool')

    result = IntegerArray(values, mask)
    expected = integer_array([1, 2, 3, np.nan], dtype='int64')
    tm.assert_extension_array_equal(result, expected)

    with pytest.raises(TypeError):
        IntegerArray(values.tolist(), mask)

    with pytest.raises(TypeError):
        IntegerArray(values, mask.tolist())

    with pytest.raises(TypeError):
        IntegerArray(values.astype(float), mask)

    with pytest.raises(TypeError):
        IntegerArray(values)
Example #5
Source File: test_datetimes.py    From recruit with Apache License 2.0 6 votes vote down vote up 
def test_fillna_preserves_tz(self, method):
        dti = pd.date_range('2000-01-01', periods=5, freq='D', tz='US/Central')
        arr = DatetimeArray(dti, copy=True)
        arr[2] = pd.NaT

        fill_val = dti[1] if method == 'pad' else dti[3]
        expected = DatetimeArray._from_sequence(
            [dti[0], dti[1], fill_val, dti[3], dti[4]],
            freq=None, tz='US/Central'
        )

        result = arr.fillna(method=method)
        tm.assert_extension_array_equal(result, expected)

        # assert that arr and dti were not modified in-place
        assert arr[2] is pd.NaT
        assert dti[2] == pd.Timestamp('2000-01-03', tz='US/Central')
Example #6
Source File: test_integer.py    From coffeegrindsize with MIT License 6 votes vote down vote up 
def test_integer_array_constructor():
    values = np.array([1, 2, 3, 4], dtype='int64')
    mask = np.array([False, False, False, True], dtype='bool')

    result = IntegerArray(values, mask)
    expected = integer_array([1, 2, 3, np.nan], dtype='int64')
    tm.assert_extension_array_equal(result, expected)

    with pytest.raises(TypeError):
        IntegerArray(values.tolist(), mask)

    with pytest.raises(TypeError):
        IntegerArray(values, mask.tolist())

    with pytest.raises(TypeError):
        IntegerArray(values.astype(float), mask)

    with pytest.raises(TypeError):
        IntegerArray(values)
Example #7
Source File: test_numpy.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_ufunc():
    arr = PandasArray(np.array([-1.0, 0.0, 1.0]))
    result = np.abs(arr)
    expected = PandasArray(np.abs(arr._ndarray))
    tm.assert_extension_array_equal(result, expected)

    r1, r2 = np.divmod(arr, np.add(arr, 2))
    e1, e2 = np.divmod(arr._ndarray, np.add(arr._ndarray, 2))
    e1 = PandasArray(e1)
    e2 = PandasArray(e2)
    tm.assert_extension_array_equal(r1, e1)
    tm.assert_extension_array_equal(r2, e2)
Example #8
Source File: test_integer.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_to_integer_array(values, to_dtype, result_dtype):
    # convert existing arrays to IntegerArrays
    result = integer_array(values, dtype=to_dtype)
    assert result.dtype == result_dtype()
    expected = integer_array(values, dtype=result_dtype())
    tm.assert_extension_array_equal(result, expected)
Example #9
Source File: test_interval.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_set_closed(self, closed, new_closed):
        # GH 21670
        array = IntervalArray.from_breaks(range(10), closed=closed)
        result = array.set_closed(new_closed)
        expected = IntervalArray.from_breaks(range(10), closed=new_closed)
        tm.assert_extension_array_equal(result, expected)
Example #10
Source File: test_interval.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_set_na(self, left_right_dtypes):
        left, right = left_right_dtypes
        result = IntervalArray.from_arrays(left, right)
        result[0] = np.nan

        expected_left = Index([left._na_value] + list(left[1:]))
        expected_right = Index([right._na_value] + list(right[1:]))
        expected = IntervalArray.from_arrays(expected_left, expected_right)

        tm.assert_extension_array_equal(result, expected)
Example #11
Source File: test_base.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_numpy_array(arr):
    ser = pd.Series(arr)
    result = ser.array
    expected = PandasArray(arr)
    tm.assert_extension_array_equal(result, expected)
Example #12
Source File: test_algos.py    From predictive-maintenance-using-machine-learning with Apache License 2.0 5 votes vote down vote up 
def test_datetime64tz_aware(self):
        # GH 15939

        result = Series(
            Index([Timestamp('20160101', tz='US/Eastern'),
                   Timestamp('20160101', tz='US/Eastern')])).unique()
        expected = DatetimeArray._from_sequence(np.array([
            Timestamp('2016-01-01 00:00:00-0500', tz="US/Eastern")
        ]))
        tm.assert_extension_array_equal(result, expected)

        result = Index([Timestamp('20160101', tz='US/Eastern'),
                        Timestamp('20160101', tz='US/Eastern')]).unique()
        expected = DatetimeIndex(['2016-01-01 00:00:00'],
                                 dtype='datetime64[ns, US/Eastern]', freq=None)
        tm.assert_index_equal(result, expected)

        result = pd.unique(
            Series(Index([Timestamp('20160101', tz='US/Eastern'),
                          Timestamp('20160101', tz='US/Eastern')])))
        expected = DatetimeArray._from_sequence(np.array([
            Timestamp('2016-01-01', tz="US/Eastern"),
        ]))
        tm.assert_extension_array_equal(result, expected)

        result = pd.unique(Index([Timestamp('20160101', tz='US/Eastern'),
                                  Timestamp('20160101', tz='US/Eastern')]))
        expected = DatetimeIndex(['2016-01-01 00:00:00'],
                                 dtype='datetime64[ns, US/Eastern]', freq=None)
        tm.assert_index_equal(result, expected)
Example #13
Source File: groupby.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_grouping_grouper(self, data_for_grouping):
        df = pd.DataFrame({
            "A": ["B", "B", None, None, "A", "A", "B", "C"],
            "B": data_for_grouping
        })
        gr1 = df.groupby("A").grouper.groupings[0]
        gr2 = df.groupby("B").grouper.groupings[0]

        tm.assert_numpy_array_equal(gr1.grouper, df.A.values)
        tm.assert_extension_array_equal(gr2.grouper, data_for_grouping)
Example #14
Source File: test_decimal.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_take_na_value_other_decimal(self):
        arr = DecimalArray([decimal.Decimal('1.0'),
                            decimal.Decimal('2.0')])
        result = arr.take([0, -1], allow_fill=True,
                          fill_value=decimal.Decimal('-1.0'))
        expected = DecimalArray([decimal.Decimal('1.0'),
                                 decimal.Decimal('-1.0')])
        self.assert_extension_array_equal(result, expected)
Example #15
Source File: test_decimal.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_divmod_array(reverse, expected_div, expected_mod):
    # https://github.com/pandas-dev/pandas/issues/22930
    arr = to_decimal([1, 2, 3, 4])
    if reverse:
        div, mod = divmod(2, arr)
    else:
        div, mod = divmod(arr, 2)
    expected_div = to_decimal(expected_div)
    expected_mod = to_decimal(expected_mod)

    tm.assert_extension_array_equal(div, expected_div)
    tm.assert_extension_array_equal(mod, expected_mod)
Example #16
Source File: test_integer.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_pow(self):
        # https://github.com/pandas-dev/pandas/issues/22022
        a = integer_array([1, np.nan, np.nan, 1])
        b = integer_array([1, np.nan, 1, np.nan])
        result = a ** b
        expected = pd.core.arrays.integer_array([1, np.nan, np.nan, 1])
        tm.assert_extension_array_equal(result, expected)
Example #17
Source File: test_integer.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_to_integer_array_float():
    result = integer_array([1., 2.])
    expected = integer_array([1, 2])
    tm.assert_extension_array_equal(result, expected)

    with pytest.raises(TypeError, match="cannot safely cast non-equivalent"):
        integer_array([1.5, 2.])

    # for float dtypes, the itemsize is not preserved
    result = integer_array(np.array([1., 2.], dtype='float32'))
    assert result.dtype == Int64Dtype()
Example #18
Source File: test_integer.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_to_integer_array(values, to_dtype, result_dtype):
    # convert existing arrays to IntegerArrays
    result = integer_array(values, dtype=to_dtype)
    assert result.dtype == result_dtype()
    expected = integer_array(values, dtype=result_dtype())
    tm.assert_extension_array_equal(result, expected)
Example #19
Source File: test_array.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_array_inference_fails(data):
    result = pd.array(data)
    expected = PandasArray(np.array(data, dtype=object))
    tm.assert_extension_array_equal(result, expected)
Example #20
Source File: test_numpy.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_from_sequence_dtype():
    arr = np.array([1, 2, 3], dtype='int64')
    result = PandasArray._from_sequence(arr, dtype='uint64')
    expected = PandasArray(np.array([1, 2, 3], dtype='uint64'))
    tm.assert_extension_array_equal(result, expected)
Example #21
Source File: test_numpy.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_ufunc():
    arr = PandasArray(np.array([-1.0, 0.0, 1.0]))
    result = np.abs(arr)
    expected = PandasArray(np.abs(arr._ndarray))
    tm.assert_extension_array_equal(result, expected)

    r1, r2 = np.divmod(arr, np.add(arr, 2))
    e1, e2 = np.divmod(arr._ndarray, np.add(arr._ndarray, 2))
    e1 = PandasArray(e1)
    e2 = PandasArray(e2)
    tm.assert_extension_array_equal(r1, e1)
    tm.assert_extension_array_equal(r2, e2)
Example #22
Source File: test_interval.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_set_closed(self, closed, new_closed):
        # GH 21670
        array = IntervalArray.from_breaks(range(10), closed=closed)
        result = array.set_closed(new_closed)
        expected = IntervalArray.from_breaks(range(10), closed=new_closed)
        tm.assert_extension_array_equal(result, expected)
Example #23
Source File: test_interval.py    From coffeegrindsize with MIT License 5 votes vote down vote up 
def test_set_na(self, left_right_dtypes):
        left, right = left_right_dtypes
        result = IntervalArray.from_arrays(left, right)
        result[0] = np.nan

        expected_left = Index([left._na_value] + list(left[1:]))
        expected_right = Index([right._na_value] + list(right[1:]))
        expected = IntervalArray.from_arrays(expected_left, expected_right)

        tm.assert_extension_array_equal(result, expected)
Example #24
Source File: groupby.py    From twitter-stock-recommendation with MIT License 5 votes vote down vote up 
def test_grouping_grouper(self, data_for_grouping):
        df = pd.DataFrame({
            "A": ["B", "B", None, None, "A", "A", "B", "C"],
            "B": data_for_grouping
        })
        gr1 = df.groupby("A").grouper.groupings[0]
        gr2 = df.groupby("B").grouper.groupings[0]

        tm.assert_numpy_array_equal(gr1.grouper, df.A.values)
        tm.assert_extension_array_equal(gr2.grouper, data_for_grouping)
Example #25
Source File: test_base.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_numpy_array(arr):
    ser = pd.Series(arr)
    result = ser.array
    expected = PandasArray(arr)
    tm.assert_extension_array_equal(result, expected)
Example #26
Source File: test_decimal.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_take_na_value_other_decimal(self):
        arr = DecimalArray([decimal.Decimal('1.0'),
                            decimal.Decimal('2.0')])
        result = arr.take([0, -1], allow_fill=True,
                          fill_value=decimal.Decimal('-1.0'))
        expected = DecimalArray([decimal.Decimal('1.0'),
                                 decimal.Decimal('-1.0')])
        self.assert_extension_array_equal(result, expected)
Example #27
Source File: test_decimal.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_divmod_array(reverse, expected_div, expected_mod):
    # https://github.com/pandas-dev/pandas/issues/22930
    arr = to_decimal([1, 2, 3, 4])
    if reverse:
        div, mod = divmod(2, arr)
    else:
        div, mod = divmod(arr, 2)
    expected_div = to_decimal(expected_div)
    expected_mod = to_decimal(expected_mod)

    tm.assert_extension_array_equal(div, expected_div)
    tm.assert_extension_array_equal(mod, expected_mod)
Example #28
Source File: test_integer.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_pow(self):
        # https://github.com/pandas-dev/pandas/issues/22022
        a = integer_array([1, np.nan, np.nan, 1])
        b = integer_array([1, np.nan, 1, np.nan])
        result = a ** b
        expected = pd.core.arrays.integer_array([1, np.nan, np.nan, 1])
        tm.assert_extension_array_equal(result, expected)
Example #29
Source File: test_integer.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_to_integer_array_float():
    result = integer_array([1., 2.])
    expected = integer_array([1, 2])
    tm.assert_extension_array_equal(result, expected)

    with pytest.raises(TypeError, match="cannot safely cast non-equivalent"):
        integer_array([1.5, 2.])

    # for float dtypes, the itemsize is not preserved
    result = integer_array(np.array([1., 2.], dtype='float32'))
    assert result.dtype == Int64Dtype()
Example #30
Source File: test_integer.py    From recruit with Apache License 2.0 5 votes vote down vote up 
def test_to_integer_array(values, to_dtype, result_dtype):
    # convert existing arrays to IntegerArrays
    result = integer_array(values, dtype=to_dtype)
    assert result.dtype == result_dtype()
    expected = integer_array(values, dtype=result_dtype())
    tm.assert_extension_array_equal(result, expected)