1import numpy as np
2import pytest
3
4import pandas as pd
5import pandas._testing as tm
6
7
8@pytest.mark.parametrize(
9    "ufunc", [np.add, np.logical_or, np.logical_and, np.logical_xor]
10)
11def test_ufuncs_binary(ufunc):
12    # two BooleanArrays
13    a = pd.array([True, False, None], dtype="boolean")
14    result = ufunc(a, a)
15    expected = pd.array(ufunc(a._data, a._data), dtype="boolean")
16    expected[a._mask] = np.nan
17    tm.assert_extension_array_equal(result, expected)
18
19    s = pd.Series(a)
20    result = ufunc(s, a)
21    expected = pd.Series(ufunc(a._data, a._data), dtype="boolean")
22    expected[a._mask] = np.nan
23    tm.assert_series_equal(result, expected)
24
25    # Boolean with numpy array
26    arr = np.array([True, True, False])
27    result = ufunc(a, arr)
28    expected = pd.array(ufunc(a._data, arr), dtype="boolean")
29    expected[a._mask] = np.nan
30    tm.assert_extension_array_equal(result, expected)
31
32    result = ufunc(arr, a)
33    expected = pd.array(ufunc(arr, a._data), dtype="boolean")
34    expected[a._mask] = np.nan
35    tm.assert_extension_array_equal(result, expected)
36
37    # BooleanArray with scalar
38    result = ufunc(a, True)
39    expected = pd.array(ufunc(a._data, True), dtype="boolean")
40    expected[a._mask] = np.nan
41    tm.assert_extension_array_equal(result, expected)
42
43    result = ufunc(True, a)
44    expected = pd.array(ufunc(True, a._data), dtype="boolean")
45    expected[a._mask] = np.nan
46    tm.assert_extension_array_equal(result, expected)
47
48    # not handled types
49    msg = r"operand type\(s\) all returned NotImplemented from __array_ufunc__"
50    with pytest.raises(TypeError, match=msg):
51        ufunc(a, "test")
52
53
54@pytest.mark.parametrize("ufunc", [np.logical_not])
55def test_ufuncs_unary(ufunc):
56    a = pd.array([True, False, None], dtype="boolean")
57    result = ufunc(a)
58    expected = pd.array(ufunc(a._data), dtype="boolean")
59    expected[a._mask] = np.nan
60    tm.assert_extension_array_equal(result, expected)
61
62    s = pd.Series(a)
63    result = ufunc(s)
64    expected = pd.Series(ufunc(a._data), dtype="boolean")
65    expected[a._mask] = np.nan
66    tm.assert_series_equal(result, expected)
67
68
69@pytest.mark.parametrize("values", [[True, False], [True, None]])
70def test_ufunc_reduce_raises(values):
71    a = pd.array(values, dtype="boolean")
72    msg = "The 'reduce' method is not supported"
73    with pytest.raises(NotImplementedError, match=msg):
74        np.add.reduce(a)
75
76
77def test_value_counts_na():
78    arr = pd.array([True, False, pd.NA], dtype="boolean")
79    result = arr.value_counts(dropna=False)
80    expected = pd.Series([1, 1, 1], index=[False, True, pd.NA], dtype="Int64")
81    tm.assert_series_equal(result, expected)
82
83    result = arr.value_counts(dropna=True)
84    expected = pd.Series([1, 1], index=[False, True], dtype="Int64")
85    tm.assert_series_equal(result, expected)
86
87
88def test_value_counts_with_normalize():
89    s = pd.Series([True, False, pd.NA], dtype="boolean")
90    result = s.value_counts(normalize=True)
91    expected = pd.Series([1, 1], index=[False, True], dtype="Float64") / 2
92    tm.assert_series_equal(result, expected)
93
94
95def test_diff():
96    a = pd.array(
97        [True, True, False, False, True, None, True, None, False], dtype="boolean"
98    )
99    result = pd.core.algorithms.diff(a, 1)
100    expected = pd.array(
101        [None, False, True, False, True, None, None, None, None], dtype="boolean"
102    )
103    tm.assert_extension_array_equal(result, expected)
104
105    s = pd.Series(a)
106    result = s.diff()
107    expected = pd.Series(expected)
108    tm.assert_series_equal(result, expected)
109