test_logical_ops.py 6.0 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191
  1. import operator
  2. import re
  3. import numpy as np
  4. import pytest
  5. from pandas import (
  6. CategoricalIndex,
  7. DataFrame,
  8. Interval,
  9. Series,
  10. isnull,
  11. )
  12. import pandas._testing as tm
  13. class TestDataFrameLogicalOperators:
  14. # &, |, ^
  15. @pytest.mark.parametrize(
  16. "left, right, op, expected",
  17. [
  18. (
  19. [True, False, np.nan],
  20. [True, False, True],
  21. operator.and_,
  22. [True, False, False],
  23. ),
  24. (
  25. [True, False, True],
  26. [True, False, np.nan],
  27. operator.and_,
  28. [True, False, False],
  29. ),
  30. (
  31. [True, False, np.nan],
  32. [True, False, True],
  33. operator.or_,
  34. [True, False, False],
  35. ),
  36. (
  37. [True, False, True],
  38. [True, False, np.nan],
  39. operator.or_,
  40. [True, False, True],
  41. ),
  42. ],
  43. )
  44. def test_logical_operators_nans(self, left, right, op, expected, frame_or_series):
  45. # GH#13896
  46. result = op(frame_or_series(left), frame_or_series(right))
  47. expected = frame_or_series(expected)
  48. tm.assert_equal(result, expected)
  49. def test_logical_ops_empty_frame(self):
  50. # GH#5808
  51. # empty frames, non-mixed dtype
  52. df = DataFrame(index=[1])
  53. result = df & df
  54. tm.assert_frame_equal(result, df)
  55. result = df | df
  56. tm.assert_frame_equal(result, df)
  57. df2 = DataFrame(index=[1, 2])
  58. result = df & df2
  59. tm.assert_frame_equal(result, df2)
  60. dfa = DataFrame(index=[1], columns=["A"])
  61. result = dfa & dfa
  62. expected = DataFrame(False, index=[1], columns=["A"])
  63. tm.assert_frame_equal(result, expected)
  64. def test_logical_ops_bool_frame(self):
  65. # GH#5808
  66. df1a_bool = DataFrame(True, index=[1], columns=["A"])
  67. result = df1a_bool & df1a_bool
  68. tm.assert_frame_equal(result, df1a_bool)
  69. result = df1a_bool | df1a_bool
  70. tm.assert_frame_equal(result, df1a_bool)
  71. def test_logical_ops_int_frame(self):
  72. # GH#5808
  73. df1a_int = DataFrame(1, index=[1], columns=["A"])
  74. df1a_bool = DataFrame(True, index=[1], columns=["A"])
  75. result = df1a_int | df1a_bool
  76. tm.assert_frame_equal(result, df1a_bool)
  77. # Check that this matches Series behavior
  78. res_ser = df1a_int["A"] | df1a_bool["A"]
  79. tm.assert_series_equal(res_ser, df1a_bool["A"])
  80. def test_logical_ops_invalid(self):
  81. # GH#5808
  82. df1 = DataFrame(1.0, index=[1], columns=["A"])
  83. df2 = DataFrame(True, index=[1], columns=["A"])
  84. msg = re.escape("unsupported operand type(s) for |: 'float' and 'bool'")
  85. with pytest.raises(TypeError, match=msg):
  86. df1 | df2
  87. df1 = DataFrame("foo", index=[1], columns=["A"])
  88. df2 = DataFrame(True, index=[1], columns=["A"])
  89. msg = re.escape("unsupported operand type(s) for |: 'str' and 'bool'")
  90. with pytest.raises(TypeError, match=msg):
  91. df1 | df2
  92. def test_logical_operators(self):
  93. def _check_bin_op(op):
  94. result = op(df1, df2)
  95. expected = DataFrame(
  96. op(df1.values, df2.values), index=df1.index, columns=df1.columns
  97. )
  98. assert result.values.dtype == np.bool_
  99. tm.assert_frame_equal(result, expected)
  100. def _check_unary_op(op):
  101. result = op(df1)
  102. expected = DataFrame(op(df1.values), index=df1.index, columns=df1.columns)
  103. assert result.values.dtype == np.bool_
  104. tm.assert_frame_equal(result, expected)
  105. df1 = {
  106. "a": {"a": True, "b": False, "c": False, "d": True, "e": True},
  107. "b": {"a": False, "b": True, "c": False, "d": False, "e": False},
  108. "c": {"a": False, "b": False, "c": True, "d": False, "e": False},
  109. "d": {"a": True, "b": False, "c": False, "d": True, "e": True},
  110. "e": {"a": True, "b": False, "c": False, "d": True, "e": True},
  111. }
  112. df2 = {
  113. "a": {"a": True, "b": False, "c": True, "d": False, "e": False},
  114. "b": {"a": False, "b": True, "c": False, "d": False, "e": False},
  115. "c": {"a": True, "b": False, "c": True, "d": False, "e": False},
  116. "d": {"a": False, "b": False, "c": False, "d": True, "e": False},
  117. "e": {"a": False, "b": False, "c": False, "d": False, "e": True},
  118. }
  119. df1 = DataFrame(df1)
  120. df2 = DataFrame(df2)
  121. _check_bin_op(operator.and_)
  122. _check_bin_op(operator.or_)
  123. _check_bin_op(operator.xor)
  124. _check_unary_op(operator.inv) # TODO: belongs elsewhere
  125. def test_logical_with_nas(self):
  126. d = DataFrame({"a": [np.nan, False], "b": [True, True]})
  127. # GH4947
  128. # bool comparisons should return bool
  129. result = d["a"] | d["b"]
  130. expected = Series([False, True])
  131. tm.assert_series_equal(result, expected)
  132. # GH4604, automatic casting here
  133. result = d["a"].fillna(False) | d["b"]
  134. expected = Series([True, True])
  135. tm.assert_series_equal(result, expected)
  136. result = d["a"].fillna(False, downcast=False) | d["b"]
  137. expected = Series([True, True])
  138. tm.assert_series_equal(result, expected)
  139. def test_logical_ops_categorical_columns(self):
  140. # GH#38367
  141. intervals = [Interval(1, 2), Interval(3, 4)]
  142. data = DataFrame(
  143. [[1, np.nan], [2, np.nan]],
  144. columns=CategoricalIndex(
  145. intervals, categories=intervals + [Interval(5, 6)]
  146. ),
  147. )
  148. mask = DataFrame(
  149. [[False, False], [False, False]], columns=data.columns, dtype=bool
  150. )
  151. result = mask | isnull(data)
  152. expected = DataFrame(
  153. [[False, True], [False, True]],
  154. columns=CategoricalIndex(
  155. intervals, categories=intervals + [Interval(5, 6)]
  156. ),
  157. )
  158. tm.assert_frame_equal(result, expected)