test_boolean.py 32 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931
  1. import operator
  2. import numpy as np
  3. import pytest
  4. import pandas.util._test_decorators as td
  5. import pandas as pd
  6. import pandas._testing as tm
  7. from pandas.arrays import BooleanArray
  8. from pandas.core.arrays.boolean import coerce_to_array
  9. from pandas.tests.extension.base import BaseOpsUtil
  10. def make_data():
  11. return [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False]
  12. @pytest.fixture
  13. def dtype():
  14. return pd.BooleanDtype()
  15. @pytest.fixture
  16. def data(dtype):
  17. return pd.array(make_data(), dtype=dtype)
  18. def test_boolean_array_constructor():
  19. values = np.array([True, False, True, False], dtype="bool")
  20. mask = np.array([False, False, False, True], dtype="bool")
  21. result = BooleanArray(values, mask)
  22. expected = pd.array([True, False, True, None], dtype="boolean")
  23. tm.assert_extension_array_equal(result, expected)
  24. with pytest.raises(TypeError, match="values should be boolean numpy array"):
  25. BooleanArray(values.tolist(), mask)
  26. with pytest.raises(TypeError, match="mask should be boolean numpy array"):
  27. BooleanArray(values, mask.tolist())
  28. with pytest.raises(TypeError, match="values should be boolean numpy array"):
  29. BooleanArray(values.astype(int), mask)
  30. with pytest.raises(TypeError, match="mask should be boolean numpy array"):
  31. BooleanArray(values, None)
  32. with pytest.raises(ValueError, match="values must be a 1D array"):
  33. BooleanArray(values.reshape(1, -1), mask)
  34. with pytest.raises(ValueError, match="mask must be a 1D array"):
  35. BooleanArray(values, mask.reshape(1, -1))
  36. def test_boolean_array_constructor_copy():
  37. values = np.array([True, False, True, False], dtype="bool")
  38. mask = np.array([False, False, False, True], dtype="bool")
  39. result = BooleanArray(values, mask)
  40. assert result._data is values
  41. assert result._mask is mask
  42. result = BooleanArray(values, mask, copy=True)
  43. assert result._data is not values
  44. assert result._mask is not mask
  45. def test_to_boolean_array():
  46. expected = BooleanArray(
  47. np.array([True, False, True]), np.array([False, False, False])
  48. )
  49. result = pd.array([True, False, True], dtype="boolean")
  50. tm.assert_extension_array_equal(result, expected)
  51. result = pd.array(np.array([True, False, True]), dtype="boolean")
  52. tm.assert_extension_array_equal(result, expected)
  53. result = pd.array(np.array([True, False, True], dtype=object), dtype="boolean")
  54. tm.assert_extension_array_equal(result, expected)
  55. # with missing values
  56. expected = BooleanArray(
  57. np.array([True, False, True]), np.array([False, False, True])
  58. )
  59. result = pd.array([True, False, None], dtype="boolean")
  60. tm.assert_extension_array_equal(result, expected)
  61. result = pd.array(np.array([True, False, None], dtype=object), dtype="boolean")
  62. tm.assert_extension_array_equal(result, expected)
  63. def test_to_boolean_array_all_none():
  64. expected = BooleanArray(np.array([True, True, True]), np.array([True, True, True]))
  65. result = pd.array([None, None, None], dtype="boolean")
  66. tm.assert_extension_array_equal(result, expected)
  67. result = pd.array(np.array([None, None, None], dtype=object), dtype="boolean")
  68. tm.assert_extension_array_equal(result, expected)
  69. @pytest.mark.parametrize(
  70. "a, b",
  71. [
  72. ([True, False, None, np.nan, pd.NA], [True, False, None, None, None]),
  73. ([True, np.nan], [True, None]),
  74. ([True, pd.NA], [True, None]),
  75. ([np.nan, np.nan], [None, None]),
  76. (np.array([np.nan, np.nan], dtype=float), [None, None]),
  77. ],
  78. )
  79. def test_to_boolean_array_missing_indicators(a, b):
  80. result = pd.array(a, dtype="boolean")
  81. expected = pd.array(b, dtype="boolean")
  82. tm.assert_extension_array_equal(result, expected)
  83. @pytest.mark.parametrize(
  84. "values",
  85. [
  86. ["foo", "bar"],
  87. ["1", "2"],
  88. # "foo",
  89. [1, 2],
  90. [1.0, 2.0],
  91. pd.date_range("20130101", periods=2),
  92. np.array(["foo"]),
  93. np.array([1, 2]),
  94. np.array([1.0, 2.0]),
  95. [np.nan, {"a": 1}],
  96. ],
  97. )
  98. def test_to_boolean_array_error(values):
  99. # error in converting existing arrays to BooleanArray
  100. with pytest.raises(TypeError):
  101. pd.array(values, dtype="boolean")
  102. def test_to_boolean_array_from_integer_array():
  103. result = pd.array(np.array([1, 0, 1, 0]), dtype="boolean")
  104. expected = pd.array([True, False, True, False], dtype="boolean")
  105. tm.assert_extension_array_equal(result, expected)
  106. # with missing values
  107. result = pd.array(np.array([1, 0, 1, None]), dtype="boolean")
  108. expected = pd.array([True, False, True, None], dtype="boolean")
  109. tm.assert_extension_array_equal(result, expected)
  110. def test_to_boolean_array_from_float_array():
  111. result = pd.array(np.array([1.0, 0.0, 1.0, 0.0]), dtype="boolean")
  112. expected = pd.array([True, False, True, False], dtype="boolean")
  113. tm.assert_extension_array_equal(result, expected)
  114. # with missing values
  115. result = pd.array(np.array([1.0, 0.0, 1.0, np.nan]), dtype="boolean")
  116. expected = pd.array([True, False, True, None], dtype="boolean")
  117. tm.assert_extension_array_equal(result, expected)
  118. def test_to_boolean_array_integer_like():
  119. # integers of 0's and 1's
  120. result = pd.array([1, 0, 1, 0], dtype="boolean")
  121. expected = pd.array([True, False, True, False], dtype="boolean")
  122. tm.assert_extension_array_equal(result, expected)
  123. # with missing values
  124. result = pd.array([1, 0, 1, None], dtype="boolean")
  125. expected = pd.array([True, False, True, None], dtype="boolean")
  126. tm.assert_extension_array_equal(result, expected)
  127. def test_coerce_to_array():
  128. # TODO this is currently not public API
  129. values = np.array([True, False, True, False], dtype="bool")
  130. mask = np.array([False, False, False, True], dtype="bool")
  131. result = BooleanArray(*coerce_to_array(values, mask=mask))
  132. expected = BooleanArray(values, mask)
  133. tm.assert_extension_array_equal(result, expected)
  134. assert result._data is values
  135. assert result._mask is mask
  136. result = BooleanArray(*coerce_to_array(values, mask=mask, copy=True))
  137. expected = BooleanArray(values, mask)
  138. tm.assert_extension_array_equal(result, expected)
  139. assert result._data is not values
  140. assert result._mask is not mask
  141. # mixed missing from values and mask
  142. values = [True, False, None, False]
  143. mask = np.array([False, False, False, True], dtype="bool")
  144. result = BooleanArray(*coerce_to_array(values, mask=mask))
  145. expected = BooleanArray(
  146. np.array([True, False, True, True]), np.array([False, False, True, True])
  147. )
  148. tm.assert_extension_array_equal(result, expected)
  149. result = BooleanArray(*coerce_to_array(np.array(values, dtype=object), mask=mask))
  150. tm.assert_extension_array_equal(result, expected)
  151. result = BooleanArray(*coerce_to_array(values, mask=mask.tolist()))
  152. tm.assert_extension_array_equal(result, expected)
  153. # raise errors for wrong dimension
  154. values = np.array([True, False, True, False], dtype="bool")
  155. mask = np.array([False, False, False, True], dtype="bool")
  156. with pytest.raises(ValueError, match="values must be a 1D list-like"):
  157. coerce_to_array(values.reshape(1, -1))
  158. with pytest.raises(ValueError, match="mask must be a 1D list-like"):
  159. coerce_to_array(values, mask=mask.reshape(1, -1))
  160. def test_coerce_to_array_from_boolean_array():
  161. # passing BooleanArray to coerce_to_array
  162. values = np.array([True, False, True, False], dtype="bool")
  163. mask = np.array([False, False, False, True], dtype="bool")
  164. arr = BooleanArray(values, mask)
  165. result = BooleanArray(*coerce_to_array(arr))
  166. tm.assert_extension_array_equal(result, arr)
  167. # no copy
  168. assert result._data is arr._data
  169. assert result._mask is arr._mask
  170. result = BooleanArray(*coerce_to_array(arr), copy=True)
  171. tm.assert_extension_array_equal(result, arr)
  172. assert result._data is not arr._data
  173. assert result._mask is not arr._mask
  174. with pytest.raises(ValueError, match="cannot pass mask for BooleanArray input"):
  175. coerce_to_array(arr, mask=mask)
  176. def test_coerce_to_numpy_array():
  177. # with missing values -> object dtype
  178. arr = pd.array([True, False, None], dtype="boolean")
  179. result = np.array(arr)
  180. expected = np.array([True, False, pd.NA], dtype="object")
  181. tm.assert_numpy_array_equal(result, expected)
  182. # also with no missing values -> object dtype
  183. arr = pd.array([True, False, True], dtype="boolean")
  184. result = np.array(arr)
  185. expected = np.array([True, False, True], dtype="object")
  186. tm.assert_numpy_array_equal(result, expected)
  187. # force bool dtype
  188. result = np.array(arr, dtype="bool")
  189. expected = np.array([True, False, True], dtype="bool")
  190. tm.assert_numpy_array_equal(result, expected)
  191. # with missing values will raise error
  192. arr = pd.array([True, False, None], dtype="boolean")
  193. with pytest.raises(ValueError):
  194. np.array(arr, dtype="bool")
  195. def test_to_boolean_array_from_strings():
  196. result = BooleanArray._from_sequence_of_strings(
  197. np.array(["True", "False", np.nan], dtype=object)
  198. )
  199. expected = BooleanArray(
  200. np.array([True, False, False]), np.array([False, False, True])
  201. )
  202. tm.assert_extension_array_equal(result, expected)
  203. def test_to_boolean_array_from_strings_invalid_string():
  204. with pytest.raises(ValueError, match="cannot be cast"):
  205. BooleanArray._from_sequence_of_strings(["donkey"])
  206. def test_repr():
  207. df = pd.DataFrame({"A": pd.array([True, False, None], dtype="boolean")})
  208. expected = " A\n0 True\n1 False\n2 <NA>"
  209. assert repr(df) == expected
  210. expected = "0 True\n1 False\n2 <NA>\nName: A, dtype: boolean"
  211. assert repr(df.A) == expected
  212. expected = "<BooleanArray>\n[True, False, <NA>]\nLength: 3, dtype: boolean"
  213. assert repr(df.A.array) == expected
  214. @pytest.mark.parametrize("box", [True, False], ids=["series", "array"])
  215. def test_to_numpy(box):
  216. con = pd.Series if box else pd.array
  217. # default (with or without missing values) -> object dtype
  218. arr = con([True, False, True], dtype="boolean")
  219. result = arr.to_numpy()
  220. expected = np.array([True, False, True], dtype="object")
  221. tm.assert_numpy_array_equal(result, expected)
  222. arr = con([True, False, None], dtype="boolean")
  223. result = arr.to_numpy()
  224. expected = np.array([True, False, pd.NA], dtype="object")
  225. tm.assert_numpy_array_equal(result, expected)
  226. arr = con([True, False, None], dtype="boolean")
  227. result = arr.to_numpy(dtype="str")
  228. expected = np.array([True, False, pd.NA], dtype="<U5")
  229. tm.assert_numpy_array_equal(result, expected)
  230. # no missing values -> can convert to bool, otherwise raises
  231. arr = con([True, False, True], dtype="boolean")
  232. result = arr.to_numpy(dtype="bool")
  233. expected = np.array([True, False, True], dtype="bool")
  234. tm.assert_numpy_array_equal(result, expected)
  235. arr = con([True, False, None], dtype="boolean")
  236. with pytest.raises(ValueError, match="cannot convert to 'bool'-dtype"):
  237. result = arr.to_numpy(dtype="bool")
  238. # specify dtype and na_value
  239. arr = con([True, False, None], dtype="boolean")
  240. result = arr.to_numpy(dtype=object, na_value=None)
  241. expected = np.array([True, False, None], dtype="object")
  242. tm.assert_numpy_array_equal(result, expected)
  243. result = arr.to_numpy(dtype=bool, na_value=False)
  244. expected = np.array([True, False, False], dtype="bool")
  245. tm.assert_numpy_array_equal(result, expected)
  246. result = arr.to_numpy(dtype="int64", na_value=-99)
  247. expected = np.array([1, 0, -99], dtype="int64")
  248. tm.assert_numpy_array_equal(result, expected)
  249. result = arr.to_numpy(dtype="float64", na_value=np.nan)
  250. expected = np.array([1, 0, np.nan], dtype="float64")
  251. tm.assert_numpy_array_equal(result, expected)
  252. # converting to int or float without specifying na_value raises
  253. with pytest.raises(ValueError, match="cannot convert to 'int64'-dtype"):
  254. arr.to_numpy(dtype="int64")
  255. with pytest.raises(ValueError, match="cannot convert to 'float64'-dtype"):
  256. arr.to_numpy(dtype="float64")
  257. def test_to_numpy_copy():
  258. # to_numpy can be zero-copy if no missing values
  259. arr = pd.array([True, False, True], dtype="boolean")
  260. result = arr.to_numpy(dtype=bool)
  261. result[0] = False
  262. tm.assert_extension_array_equal(
  263. arr, pd.array([False, False, True], dtype="boolean")
  264. )
  265. arr = pd.array([True, False, True], dtype="boolean")
  266. result = arr.to_numpy(dtype=bool, copy=True)
  267. result[0] = False
  268. tm.assert_extension_array_equal(arr, pd.array([True, False, True], dtype="boolean"))
  269. def test_astype():
  270. # with missing values
  271. arr = pd.array([True, False, None], dtype="boolean")
  272. with pytest.raises(ValueError, match="cannot convert NA to integer"):
  273. arr.astype("int64")
  274. with pytest.raises(ValueError, match="cannot convert float NaN to"):
  275. arr.astype("bool")
  276. result = arr.astype("float64")
  277. expected = np.array([1, 0, np.nan], dtype="float64")
  278. tm.assert_numpy_array_equal(result, expected)
  279. result = arr.astype("str")
  280. expected = np.array(["True", "False", "<NA>"], dtype="object")
  281. tm.assert_numpy_array_equal(result, expected)
  282. # no missing values
  283. arr = pd.array([True, False, True], dtype="boolean")
  284. result = arr.astype("int64")
  285. expected = np.array([1, 0, 1], dtype="int64")
  286. tm.assert_numpy_array_equal(result, expected)
  287. result = arr.astype("bool")
  288. expected = np.array([True, False, True], dtype="bool")
  289. tm.assert_numpy_array_equal(result, expected)
  290. def test_astype_to_boolean_array():
  291. # astype to BooleanArray
  292. arr = pd.array([True, False, None], dtype="boolean")
  293. result = arr.astype("boolean")
  294. tm.assert_extension_array_equal(result, arr)
  295. result = arr.astype(pd.BooleanDtype())
  296. tm.assert_extension_array_equal(result, arr)
  297. def test_astype_to_integer_array():
  298. # astype to IntegerArray
  299. arr = pd.array([True, False, None], dtype="boolean")
  300. result = arr.astype("Int64")
  301. expected = pd.array([1, 0, None], dtype="Int64")
  302. tm.assert_extension_array_equal(result, expected)
  303. @pytest.mark.parametrize("na", [None, np.nan, pd.NA])
  304. def test_setitem_missing_values(na):
  305. arr = pd.array([True, False, None], dtype="boolean")
  306. expected = pd.array([True, None, None], dtype="boolean")
  307. arr[1] = na
  308. tm.assert_extension_array_equal(arr, expected)
  309. @pytest.mark.parametrize(
  310. "ufunc", [np.add, np.logical_or, np.logical_and, np.logical_xor]
  311. )
  312. def test_ufuncs_binary(ufunc):
  313. # two BooleanArrays
  314. a = pd.array([True, False, None], dtype="boolean")
  315. result = ufunc(a, a)
  316. expected = pd.array(ufunc(a._data, a._data), dtype="boolean")
  317. expected[a._mask] = np.nan
  318. tm.assert_extension_array_equal(result, expected)
  319. s = pd.Series(a)
  320. result = ufunc(s, a)
  321. expected = pd.Series(ufunc(a._data, a._data), dtype="boolean")
  322. expected[a._mask] = np.nan
  323. tm.assert_series_equal(result, expected)
  324. # Boolean with numpy array
  325. arr = np.array([True, True, False])
  326. result = ufunc(a, arr)
  327. expected = pd.array(ufunc(a._data, arr), dtype="boolean")
  328. expected[a._mask] = np.nan
  329. tm.assert_extension_array_equal(result, expected)
  330. result = ufunc(arr, a)
  331. expected = pd.array(ufunc(arr, a._data), dtype="boolean")
  332. expected[a._mask] = np.nan
  333. tm.assert_extension_array_equal(result, expected)
  334. # BooleanArray with scalar
  335. result = ufunc(a, True)
  336. expected = pd.array(ufunc(a._data, True), dtype="boolean")
  337. expected[a._mask] = np.nan
  338. tm.assert_extension_array_equal(result, expected)
  339. result = ufunc(True, a)
  340. expected = pd.array(ufunc(True, a._data), dtype="boolean")
  341. expected[a._mask] = np.nan
  342. tm.assert_extension_array_equal(result, expected)
  343. # not handled types
  344. with pytest.raises(TypeError):
  345. ufunc(a, "test")
  346. @pytest.mark.parametrize("ufunc", [np.logical_not])
  347. def test_ufuncs_unary(ufunc):
  348. a = pd.array([True, False, None], dtype="boolean")
  349. result = ufunc(a)
  350. expected = pd.array(ufunc(a._data), dtype="boolean")
  351. expected[a._mask] = np.nan
  352. tm.assert_extension_array_equal(result, expected)
  353. s = pd.Series(a)
  354. result = ufunc(s)
  355. expected = pd.Series(ufunc(a._data), dtype="boolean")
  356. expected[a._mask] = np.nan
  357. tm.assert_series_equal(result, expected)
  358. @pytest.mark.parametrize("values", [[True, False], [True, None]])
  359. def test_ufunc_reduce_raises(values):
  360. a = pd.array(values, dtype="boolean")
  361. with pytest.raises(NotImplementedError):
  362. np.add.reduce(a)
  363. class TestUnaryOps:
  364. def test_invert(self):
  365. a = pd.array([True, False, None], dtype="boolean")
  366. expected = pd.array([False, True, None], dtype="boolean")
  367. tm.assert_extension_array_equal(~a, expected)
  368. expected = pd.Series(expected, index=["a", "b", "c"], name="name")
  369. result = ~pd.Series(a, index=["a", "b", "c"], name="name")
  370. tm.assert_series_equal(result, expected)
  371. df = pd.DataFrame({"A": a, "B": [True, False, False]}, index=["a", "b", "c"])
  372. result = ~df
  373. expected = pd.DataFrame(
  374. {"A": expected, "B": [False, True, True]}, index=["a", "b", "c"]
  375. )
  376. tm.assert_frame_equal(result, expected)
  377. class TestLogicalOps(BaseOpsUtil):
  378. def test_numpy_scalars_ok(self, all_logical_operators):
  379. a = pd.array([True, False, None], dtype="boolean")
  380. op = getattr(a, all_logical_operators)
  381. tm.assert_extension_array_equal(op(True), op(np.bool(True)))
  382. tm.assert_extension_array_equal(op(False), op(np.bool(False)))
  383. def get_op_from_name(self, op_name):
  384. short_opname = op_name.strip("_")
  385. short_opname = short_opname if "xor" in short_opname else short_opname + "_"
  386. try:
  387. op = getattr(operator, short_opname)
  388. except AttributeError:
  389. # Assume it is the reverse operator
  390. rop = getattr(operator, short_opname[1:])
  391. op = lambda x, y: rop(y, x)
  392. return op
  393. def test_empty_ok(self, all_logical_operators):
  394. a = pd.array([], dtype="boolean")
  395. op_name = all_logical_operators
  396. result = getattr(a, op_name)(True)
  397. tm.assert_extension_array_equal(a, result)
  398. result = getattr(a, op_name)(False)
  399. tm.assert_extension_array_equal(a, result)
  400. # TODO: pd.NA
  401. # result = getattr(a, op_name)(pd.NA)
  402. # tm.assert_extension_array_equal(a, result)
  403. def test_logical_length_mismatch_raises(self, all_logical_operators):
  404. op_name = all_logical_operators
  405. a = pd.array([True, False, None], dtype="boolean")
  406. msg = "Lengths must match to compare"
  407. with pytest.raises(ValueError, match=msg):
  408. getattr(a, op_name)([True, False])
  409. with pytest.raises(ValueError, match=msg):
  410. getattr(a, op_name)(np.array([True, False]))
  411. with pytest.raises(ValueError, match=msg):
  412. getattr(a, op_name)(pd.array([True, False], dtype="boolean"))
  413. def test_logical_nan_raises(self, all_logical_operators):
  414. op_name = all_logical_operators
  415. a = pd.array([True, False, None], dtype="boolean")
  416. msg = "Got float instead"
  417. with pytest.raises(TypeError, match=msg):
  418. getattr(a, op_name)(np.nan)
  419. @pytest.mark.parametrize("other", ["a", 1])
  420. def test_non_bool_or_na_other_raises(self, other, all_logical_operators):
  421. a = pd.array([True, False], dtype="boolean")
  422. with pytest.raises(TypeError, match=str(type(other).__name__)):
  423. getattr(a, all_logical_operators)(other)
  424. def test_kleene_or(self):
  425. # A clear test of behavior.
  426. a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  427. b = pd.array([True, False, None] * 3, dtype="boolean")
  428. result = a | b
  429. expected = pd.array(
  430. [True, True, True, True, False, None, True, None, None], dtype="boolean"
  431. )
  432. tm.assert_extension_array_equal(result, expected)
  433. result = b | a
  434. tm.assert_extension_array_equal(result, expected)
  435. # ensure we haven't mutated anything inplace
  436. tm.assert_extension_array_equal(
  437. a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  438. )
  439. tm.assert_extension_array_equal(
  440. b, pd.array([True, False, None] * 3, dtype="boolean")
  441. )
  442. @pytest.mark.parametrize(
  443. "other, expected",
  444. [
  445. (pd.NA, [True, None, None]),
  446. (True, [True, True, True]),
  447. (np.bool_(True), [True, True, True]),
  448. (False, [True, False, None]),
  449. (np.bool_(False), [True, False, None]),
  450. ],
  451. )
  452. def test_kleene_or_scalar(self, other, expected):
  453. # TODO: test True & False
  454. a = pd.array([True, False, None], dtype="boolean")
  455. result = a | other
  456. expected = pd.array(expected, dtype="boolean")
  457. tm.assert_extension_array_equal(result, expected)
  458. result = other | a
  459. tm.assert_extension_array_equal(result, expected)
  460. # ensure we haven't mutated anything inplace
  461. tm.assert_extension_array_equal(
  462. a, pd.array([True, False, None], dtype="boolean")
  463. )
  464. def test_kleene_and(self):
  465. # A clear test of behavior.
  466. a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  467. b = pd.array([True, False, None] * 3, dtype="boolean")
  468. result = a & b
  469. expected = pd.array(
  470. [True, False, None, False, False, False, None, False, None], dtype="boolean"
  471. )
  472. tm.assert_extension_array_equal(result, expected)
  473. result = b & a
  474. tm.assert_extension_array_equal(result, expected)
  475. # ensure we haven't mutated anything inplace
  476. tm.assert_extension_array_equal(
  477. a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  478. )
  479. tm.assert_extension_array_equal(
  480. b, pd.array([True, False, None] * 3, dtype="boolean")
  481. )
  482. @pytest.mark.parametrize(
  483. "other, expected",
  484. [
  485. (pd.NA, [None, False, None]),
  486. (True, [True, False, None]),
  487. (False, [False, False, False]),
  488. (np.bool_(True), [True, False, None]),
  489. (np.bool_(False), [False, False, False]),
  490. ],
  491. )
  492. def test_kleene_and_scalar(self, other, expected):
  493. a = pd.array([True, False, None], dtype="boolean")
  494. result = a & other
  495. expected = pd.array(expected, dtype="boolean")
  496. tm.assert_extension_array_equal(result, expected)
  497. result = other & a
  498. tm.assert_extension_array_equal(result, expected)
  499. # ensure we haven't mutated anything inplace
  500. tm.assert_extension_array_equal(
  501. a, pd.array([True, False, None], dtype="boolean")
  502. )
  503. def test_kleene_xor(self):
  504. a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  505. b = pd.array([True, False, None] * 3, dtype="boolean")
  506. result = a ^ b
  507. expected = pd.array(
  508. [False, True, None, True, False, None, None, None, None], dtype="boolean"
  509. )
  510. tm.assert_extension_array_equal(result, expected)
  511. result = b ^ a
  512. tm.assert_extension_array_equal(result, expected)
  513. # ensure we haven't mutated anything inplace
  514. tm.assert_extension_array_equal(
  515. a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  516. )
  517. tm.assert_extension_array_equal(
  518. b, pd.array([True, False, None] * 3, dtype="boolean")
  519. )
  520. @pytest.mark.parametrize(
  521. "other, expected",
  522. [
  523. (pd.NA, [None, None, None]),
  524. (True, [False, True, None]),
  525. (np.bool_(True), [False, True, None]),
  526. (np.bool_(False), [True, False, None]),
  527. ],
  528. )
  529. def test_kleene_xor_scalar(self, other, expected):
  530. a = pd.array([True, False, None], dtype="boolean")
  531. result = a ^ other
  532. expected = pd.array(expected, dtype="boolean")
  533. tm.assert_extension_array_equal(result, expected)
  534. result = other ^ a
  535. tm.assert_extension_array_equal(result, expected)
  536. # ensure we haven't mutated anything inplace
  537. tm.assert_extension_array_equal(
  538. a, pd.array([True, False, None], dtype="boolean")
  539. )
  540. @pytest.mark.parametrize(
  541. "other", [True, False, pd.NA, [True, False, None] * 3],
  542. )
  543. def test_no_masked_assumptions(self, other, all_logical_operators):
  544. # The logical operations should not assume that masked values are False!
  545. a = pd.arrays.BooleanArray(
  546. np.array([True, True, True, False, False, False, True, False, True]),
  547. np.array([False] * 6 + [True, True, True]),
  548. )
  549. b = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  550. if isinstance(other, list):
  551. other = pd.array(other, dtype="boolean")
  552. result = getattr(a, all_logical_operators)(other)
  553. expected = getattr(b, all_logical_operators)(other)
  554. tm.assert_extension_array_equal(result, expected)
  555. if isinstance(other, BooleanArray):
  556. other._data[other._mask] = True
  557. a._data[a._mask] = False
  558. result = getattr(a, all_logical_operators)(other)
  559. expected = getattr(b, all_logical_operators)(other)
  560. tm.assert_extension_array_equal(result, expected)
  561. class TestComparisonOps(BaseOpsUtil):
  562. def _compare_other(self, data, op_name, other):
  563. op = self.get_op_from_name(op_name)
  564. # array
  565. result = pd.Series(op(data, other))
  566. expected = pd.Series(op(data._data, other), dtype="boolean")
  567. # propagate NAs
  568. expected[data._mask] = pd.NA
  569. tm.assert_series_equal(result, expected)
  570. # series
  571. s = pd.Series(data)
  572. result = op(s, other)
  573. expected = pd.Series(data._data)
  574. expected = op(expected, other)
  575. expected = expected.astype("boolean")
  576. # propagate NAs
  577. expected[data._mask] = pd.NA
  578. tm.assert_series_equal(result, expected)
  579. def test_compare_scalar(self, data, all_compare_operators):
  580. op_name = all_compare_operators
  581. self._compare_other(data, op_name, True)
  582. def test_compare_array(self, data, all_compare_operators):
  583. op_name = all_compare_operators
  584. other = pd.array([True] * len(data), dtype="boolean")
  585. self._compare_other(data, op_name, other)
  586. other = np.array([True] * len(data))
  587. self._compare_other(data, op_name, other)
  588. other = pd.Series([True] * len(data))
  589. self._compare_other(data, op_name, other)
  590. @pytest.mark.parametrize("other", [True, False, pd.NA])
  591. def test_scalar(self, other, all_compare_operators):
  592. op = self.get_op_from_name(all_compare_operators)
  593. a = pd.array([True, False, None], dtype="boolean")
  594. result = op(a, other)
  595. if other is pd.NA:
  596. expected = pd.array([None, None, None], dtype="boolean")
  597. else:
  598. values = op(a._data, other)
  599. expected = BooleanArray(values, a._mask, copy=True)
  600. tm.assert_extension_array_equal(result, expected)
  601. # ensure we haven't mutated anything inplace
  602. result[0] = None
  603. tm.assert_extension_array_equal(
  604. a, pd.array([True, False, None], dtype="boolean")
  605. )
  606. def test_array(self, all_compare_operators):
  607. op = self.get_op_from_name(all_compare_operators)
  608. a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  609. b = pd.array([True, False, None] * 3, dtype="boolean")
  610. result = op(a, b)
  611. values = op(a._data, b._data)
  612. mask = a._mask | b._mask
  613. expected = BooleanArray(values, mask)
  614. tm.assert_extension_array_equal(result, expected)
  615. # ensure we haven't mutated anything inplace
  616. result[0] = None
  617. tm.assert_extension_array_equal(
  618. a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  619. )
  620. tm.assert_extension_array_equal(
  621. b, pd.array([True, False, None] * 3, dtype="boolean")
  622. )
  623. class TestArithmeticOps(BaseOpsUtil):
  624. def test_error(self, data, all_arithmetic_operators):
  625. # invalid ops
  626. op = all_arithmetic_operators
  627. s = pd.Series(data)
  628. ops = getattr(s, op)
  629. opa = getattr(data, op)
  630. # invalid scalars
  631. with pytest.raises(TypeError):
  632. ops("foo")
  633. with pytest.raises(TypeError):
  634. ops(pd.Timestamp("20180101"))
  635. # invalid array-likes
  636. if op not in ("__mul__", "__rmul__"):
  637. # TODO(extension) numpy's mul with object array sees booleans as numbers
  638. with pytest.raises(TypeError):
  639. ops(pd.Series("foo", index=s.index))
  640. # 2d
  641. result = opa(pd.DataFrame({"A": s}))
  642. assert result is NotImplemented
  643. with pytest.raises(NotImplementedError):
  644. opa(np.arange(len(s)).reshape(-1, len(s)))
  645. @pytest.mark.parametrize("dropna", [True, False])
  646. def test_reductions_return_types(dropna, data, all_numeric_reductions):
  647. op = all_numeric_reductions
  648. s = pd.Series(data)
  649. if dropna:
  650. s = s.dropna()
  651. if op in ("sum", "prod"):
  652. assert isinstance(getattr(s, op)(), np.int64)
  653. elif op in ("min", "max"):
  654. assert isinstance(getattr(s, op)(), np.bool_)
  655. else:
  656. # "mean", "std", "var", "median", "kurt", "skew"
  657. assert isinstance(getattr(s, op)(), np.float64)
  658. @pytest.mark.parametrize(
  659. "values, exp_any, exp_all, exp_any_noskip, exp_all_noskip",
  660. [
  661. ([True, pd.NA], True, True, True, pd.NA),
  662. ([False, pd.NA], False, False, pd.NA, False),
  663. ([pd.NA], False, True, pd.NA, pd.NA),
  664. ([], False, True, False, True),
  665. ],
  666. )
  667. def test_any_all(values, exp_any, exp_all, exp_any_noskip, exp_all_noskip):
  668. # the methods return numpy scalars
  669. exp_any = pd.NA if exp_any is pd.NA else np.bool_(exp_any)
  670. exp_all = pd.NA if exp_all is pd.NA else np.bool_(exp_all)
  671. exp_any_noskip = pd.NA if exp_any_noskip is pd.NA else np.bool_(exp_any_noskip)
  672. exp_all_noskip = pd.NA if exp_all_noskip is pd.NA else np.bool_(exp_all_noskip)
  673. for con in [pd.array, pd.Series]:
  674. a = con(values, dtype="boolean")
  675. assert a.any() is exp_any
  676. assert a.all() is exp_all
  677. assert a.any(skipna=False) is exp_any_noskip
  678. assert a.all(skipna=False) is exp_all_noskip
  679. assert np.any(a.any()) is exp_any
  680. assert np.all(a.all()) is exp_all
  681. # TODO when BooleanArray coerces to object dtype numpy array, need to do conversion
  682. # manually in the indexing code
  683. # def test_indexing_boolean_mask():
  684. # arr = pd.array([1, 2, 3, 4], dtype="Int64")
  685. # mask = pd.array([True, False, True, False], dtype="boolean")
  686. # result = arr[mask]
  687. # expected = pd.array([1, 3], dtype="Int64")
  688. # tm.assert_extension_array_equal(result, expected)
  689. # # missing values -> error
  690. # mask = pd.array([True, False, True, None], dtype="boolean")
  691. # with pytest.raises(IndexError):
  692. # result = arr[mask]
  693. @td.skip_if_no("pyarrow", min_version="0.15.0")
  694. def test_arrow_array(data):
  695. # protocol added in 0.15.0
  696. import pyarrow as pa
  697. arr = pa.array(data)
  698. # TODO use to_numpy(na_value=None) here
  699. data_object = np.array(data, dtype=object)
  700. data_object[data.isna()] = None
  701. expected = pa.array(data_object, type=pa.bool_(), from_pandas=True)
  702. assert arr.equals(expected)
  703. @td.skip_if_no("pyarrow", min_version="0.15.1.dev")
  704. def test_arrow_roundtrip():
  705. # roundtrip possible from arrow 1.0.0
  706. import pyarrow as pa
  707. data = pd.array([True, False, None], dtype="boolean")
  708. df = pd.DataFrame({"a": data})
  709. table = pa.table(df)
  710. assert table.field("a").type == "bool"
  711. result = table.to_pandas()
  712. assert isinstance(result["a"].dtype, pd.BooleanDtype)
  713. tm.assert_frame_equal(result, df)
  714. def test_value_counts_na():
  715. arr = pd.array([True, False, pd.NA], dtype="boolean")
  716. result = arr.value_counts(dropna=False)
  717. expected = pd.Series([1, 1, 1], index=[True, False, pd.NA], dtype="Int64")
  718. tm.assert_series_equal(result, expected)
  719. result = arr.value_counts(dropna=True)
  720. expected = pd.Series([1, 1], index=[True, False], dtype="Int64")
  721. tm.assert_series_equal(result, expected)
  722. def test_diff():
  723. a = pd.array(
  724. [True, True, False, False, True, None, True, None, False], dtype="boolean"
  725. )
  726. result = pd.core.algorithms.diff(a, 1)
  727. expected = pd.array(
  728. [None, False, True, False, True, None, None, None, None], dtype="boolean"
  729. )
  730. tm.assert_extension_array_equal(result, expected)
  731. s = pd.Series(a)
  732. result = s.diff()
  733. expected = pd.Series(expected)
  734. tm.assert_series_equal(result, expected)