Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 7 additions & 41 deletions pandas/tests/arrays/boolean/test_comparison.py
Original file line number Diff line number Diff line change
Expand Up @@ -4,7 +4,7 @@
import pandas as pd
import pandas._testing as tm
from pandas.arrays import BooleanArray
from pandas.tests.extension.base import BaseOpsUtil
from pandas.tests.arrays.masked_shared import ComparisonOps


@pytest.fixture
Expand All @@ -15,30 +15,12 @@ def data():
)


class TestComparisonOps(BaseOpsUtil):
def _compare_other(self, data, op_name, other):
op = self.get_op_from_name(op_name)

# array
result = pd.Series(op(data, other))
expected = pd.Series(op(data._data, other), dtype="boolean")
# propagate NAs
expected[data._mask] = pd.NA

tm.assert_series_equal(result, expected)

# series
s = pd.Series(data)
result = op(s, other)

expected = pd.Series(data._data)
expected = op(expected, other)
expected = expected.astype("boolean")
# propagate NAs
expected[data._mask] = pd.NA
@pytest.fixture
def dtype():
return pd.BooleanDtype()

tm.assert_series_equal(result, expected)

class TestComparisonOps(ComparisonOps):
def test_compare_scalar(self, data, all_compare_operators):
op_name = all_compare_operators
self._compare_other(data, op_name, True)
Expand All @@ -53,24 +35,8 @@ def test_compare_array(self, data, all_compare_operators):
self._compare_other(data, op_name, other)

@pytest.mark.parametrize("other", [True, False, pd.NA])
def test_scalar(self, other, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([True, False, None], dtype="boolean")

result = op(a, other)

if other is pd.NA:
expected = pd.array([None, None, None], dtype="boolean")
else:
values = op(a._data, other)
expected = BooleanArray(values, a._mask, copy=True)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = None
tm.assert_extension_array_equal(
a, pd.array([True, False, None], dtype="boolean")
)
def test_scalar(self, other, all_compare_operators, dtype):
ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)

def test_array(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
Expand Down
95 changes: 7 additions & 88 deletions pandas/tests/arrays/floating/test_comparison.py
Original file line number Diff line number Diff line change
@@ -1,86 +1,17 @@
import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm
from pandas.tests.extension.base import BaseOpsUtil
from pandas.tests.arrays.masked_shared import (
ComparisonOps,
NumericOps,
)


class TestComparisonOps(BaseOpsUtil):
def _compare_other(self, data, op_name, other):
op = self.get_op_from_name(op_name)

# array
result = pd.Series(op(data, other))
expected = pd.Series(op(data._data, other), dtype="boolean")

# fill the nan locations
expected[data._mask] = pd.NA

tm.assert_series_equal(result, expected)

# series
s = pd.Series(data)
result = op(s, other)

expected = op(pd.Series(data._data), other)

# fill the nan locations
expected[data._mask] = pd.NA
expected = expected.astype("boolean")

tm.assert_series_equal(result, expected)

class TestComparisonOps(NumericOps, ComparisonOps):
@pytest.mark.parametrize("other", [True, False, pd.NA, -1.0, 0.0, 1])
def test_scalar(self, other, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([1.0, 0.0, None], dtype="Float64")

result = op(a, other)

if other is pd.NA:
expected = pd.array([None, None, None], dtype="boolean")
else:
values = op(a._data, other)
expected = pd.arrays.BooleanArray(values, a._mask, copy=True)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(a, pd.array([1.0, 0.0, None], dtype="Float64"))

def test_array(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([0, 1, 2, None, None, None], dtype="Float64")
b = pd.array([0, 1, None, 0, 1, None], dtype="Float64")

result = op(a, b)
values = op(a._data, b._data)
mask = a._mask | b._mask

expected = pd.arrays.BooleanArray(values, mask)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(
a, pd.array([0, 1, 2, None, None, None], dtype="Float64")
)
tm.assert_extension_array_equal(
b, pd.array([0, 1, None, 0, 1, None], dtype="Float64")
)

def test_compare_with_booleanarray(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([True, False, None] * 3, dtype="boolean")
b = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype="Float64")
other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")
expected = op(a, other)
result = op(a, b)
tm.assert_extension_array_equal(result, expected)
expected = op(other, a)
result = op(b, a)
tm.assert_extension_array_equal(result, expected)
def test_scalar(self, other, all_compare_operators, dtype):
ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)

def test_compare_with_integerarray(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
Expand All @@ -94,18 +25,6 @@ def test_compare_with_integerarray(self, all_compare_operators):
result = op(b, a)
tm.assert_extension_array_equal(result, expected)

def test_no_shared_mask(self, data):
result = data + 1
assert np.shares_memory(result._mask, data._mask) is False

def test_compare_to_string(self, dtype):
# GH 28930
s = pd.Series([1, None], dtype=dtype)
result = s == "a"
expected = pd.Series([False, pd.NA], dtype="boolean")

self.assert_series_equal(result, expected)


def test_equals():
# GH-30652
Expand Down
97 changes: 9 additions & 88 deletions pandas/tests/arrays/integer/test_comparison.py
Original file line number Diff line number Diff line change
@@ -1,99 +1,20 @@
import numpy as np
import pytest

import pandas as pd
import pandas._testing as tm
from pandas.tests.extension.base import BaseOpsUtil
from pandas.tests.arrays.masked_shared import (
ComparisonOps,
NumericOps,
)


class TestComparisonOps(BaseOpsUtil):
def _compare_other(self, data, op_name, other):
op = self.get_op_from_name(op_name)

# array
result = pd.Series(op(data, other))
expected = pd.Series(op(data._data, other), dtype="boolean")

# fill the nan locations
expected[data._mask] = pd.NA

tm.assert_series_equal(result, expected)

# series
s = pd.Series(data)
result = op(s, other)

expected = op(pd.Series(data._data), other)

# fill the nan locations
expected[data._mask] = pd.NA
expected = expected.astype("boolean")

tm.assert_series_equal(result, expected)

class TestComparisonOps(NumericOps, ComparisonOps):
@pytest.mark.parametrize("other", [True, False, pd.NA, -1, 0, 1])
def test_scalar(self, other, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([1, 0, None], dtype="Int64")

result = op(a, other)

if other is pd.NA:
expected = pd.array([None, None, None], dtype="boolean")
else:
values = op(a._data, other)
expected = pd.arrays.BooleanArray(values, a._mask, copy=True)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(a, pd.array([1, 0, None], dtype="Int64"))

def test_array(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([0, 1, 2, None, None, None], dtype="Int64")
b = pd.array([0, 1, None, 0, 1, None], dtype="Int64")

result = op(a, b)
values = op(a._data, b._data)
mask = a._mask | b._mask

expected = pd.arrays.BooleanArray(values, mask)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(
a, pd.array([0, 1, 2, None, None, None], dtype="Int64")
)
tm.assert_extension_array_equal(
b, pd.array([0, 1, None, 0, 1, None], dtype="Int64")
)

def test_compare_with_booleanarray(self, all_compare_operators):
op = self.get_op_from_name(all_compare_operators)
a = pd.array([True, False, None] * 3, dtype="boolean")
b = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype="Int64")
other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")
expected = op(a, other)
result = op(a, b)
tm.assert_extension_array_equal(result, expected)

def test_no_shared_mask(self, data):
result = data + 1
assert np.shares_memory(result._mask, data._mask) is False

def test_compare_to_string(self, any_int_ea_dtype):
# GH 28930
s = pd.Series([1, None], dtype=any_int_ea_dtype)
result = s == "a"
expected = pd.Series([False, pd.NA], dtype="boolean")

self.assert_series_equal(result, expected)
def test_scalar(self, other, all_compare_operators, dtype):
ComparisonOps.test_scalar(self, other, all_compare_operators, dtype)

def test_compare_to_int(self, any_int_ea_dtype, all_compare_operators):
def test_compare_to_int(self, dtype, all_compare_operators):
# GH 28930
s1 = pd.Series([1, None, 3], dtype=any_int_ea_dtype)
s1 = pd.Series([1, None, 3], dtype=dtype)
s2 = pd.Series([1, None, 3], dtype="float")

method = getattr(s1, all_compare_operators)
Expand Down
106 changes: 106 additions & 0 deletions pandas/tests/arrays/masked_shared.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,106 @@
"""
Tests shared by MaskedArray subclasses.
"""
import numpy as np

import pandas as pd
import pandas._testing as tm
from pandas.tests.extension.base import BaseOpsUtil


class ComparisonOps(BaseOpsUtil):
def _compare_other(self, data, op_name, other):
op = self.get_op_from_name(op_name)

# array
result = pd.Series(op(data, other))
expected = pd.Series(op(data._data, other), dtype="boolean")

# fill the nan locations
expected[data._mask] = pd.NA

tm.assert_series_equal(result, expected)

# series
ser = pd.Series(data)
result = op(ser, other)

expected = op(pd.Series(data._data), other)

# fill the nan locations
expected[data._mask] = pd.NA
expected = expected.astype("boolean")

tm.assert_series_equal(result, expected)

# subclass will override to parametrize 'other'
def test_scalar(self, other, all_compare_operators, dtype):
op = self.get_op_from_name(all_compare_operators)
left = pd.array([1, 0, None], dtype=dtype)

result = op(left, other)

if other is pd.NA:
expected = pd.array([None, None, None], dtype="boolean")
else:
values = op(left._data, other)
expected = pd.arrays.BooleanArray(values, left._mask, copy=True)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(left, pd.array([1, 0, None], dtype=dtype))


class NumericOps:
# Shared by IntegerArray and FloatingArray, not BooleanArray

def test_no_shared_mask(self, data):
result = data + 1
assert np.shares_memory(result._mask, data._mask) is False

def test_array(self, all_compare_operators, dtype):
op = self.get_op_from_name(all_compare_operators)

left = pd.array([0, 1, 2, None, None, None], dtype=dtype)
right = pd.array([0, 1, None, 0, 1, None], dtype=dtype)

result = op(left, right)
values = op(left._data, right._data)
mask = left._mask | right._mask

expected = pd.arrays.BooleanArray(values, mask)
tm.assert_extension_array_equal(result, expected)

# ensure we haven't mutated anything inplace
result[0] = pd.NA
tm.assert_extension_array_equal(
left, pd.array([0, 1, 2, None, None, None], dtype=dtype)
)
tm.assert_extension_array_equal(
right, pd.array([0, 1, None, 0, 1, None], dtype=dtype)
)

def test_compare_with_booleanarray(self, all_compare_operators, dtype):
op = self.get_op_from_name(all_compare_operators)

left = pd.array([True, False, None] * 3, dtype="boolean")
right = pd.array([0] * 3 + [1] * 3 + [None] * 3, dtype=dtype)
other = pd.array([False] * 3 + [True] * 3 + [None] * 3, dtype="boolean")

expected = op(left, other)
result = op(left, right)
tm.assert_extension_array_equal(result, expected)

# reversed op
expected = op(other, left)
result = op(right, left)
tm.assert_extension_array_equal(result, expected)

def test_compare_to_string(self, dtype):
# GH#28930
ser = pd.Series([1, None], dtype=dtype)
result = ser == "a"
expected = pd.Series([False, pd.NA], dtype="boolean")

self.assert_series_equal(result, expected)
Loading