test_comparison.py 1.9 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960
  1. import numpy as np
  2. import pytest
  3. import pandas as pd
  4. import pandas._testing as tm
  5. from pandas.arrays import BooleanArray
  6. from pandas.tests.arrays.masked_shared import ComparisonOps
  7. @pytest.fixture
  8. def data():
  9. """Fixture returning boolean array with valid and missing data"""
  10. return pd.array(
  11. [True, False] * 4 + [np.nan] + [True, False] * 44 + [np.nan] + [True, False],
  12. dtype="boolean",
  13. )
  14. @pytest.fixture
  15. def dtype():
  16. """Fixture returning BooleanDtype"""
  17. return pd.BooleanDtype()
  18. class TestComparisonOps(ComparisonOps):
  19. def test_compare_scalar(self, data, comparison_op):
  20. self._compare_other(data, comparison_op, True)
  21. def test_compare_array(self, data, comparison_op):
  22. other = pd.array([True] * len(data), dtype="boolean")
  23. self._compare_other(data, comparison_op, other)
  24. other = np.array([True] * len(data))
  25. self._compare_other(data, comparison_op, other)
  26. other = pd.Series([True] * len(data))
  27. self._compare_other(data, comparison_op, other)
  28. @pytest.mark.parametrize("other", [True, False, pd.NA])
  29. def test_scalar(self, other, comparison_op, dtype):
  30. ComparisonOps.test_scalar(self, other, comparison_op, dtype)
  31. def test_array(self, comparison_op):
  32. op = comparison_op
  33. a = pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  34. b = pd.array([True, False, None] * 3, dtype="boolean")
  35. result = op(a, b)
  36. values = op(a._data, b._data)
  37. mask = a._mask | b._mask
  38. expected = BooleanArray(values, mask)
  39. tm.assert_extension_array_equal(result, expected)
  40. # ensure we haven't mutated anything inplace
  41. result[0] = None
  42. tm.assert_extension_array_equal(
  43. a, pd.array([True] * 3 + [False] * 3 + [None] * 3, dtype="boolean")
  44. )
  45. tm.assert_extension_array_equal(
  46. b, pd.array([True, False, None] * 3, dtype="boolean")
  47. )