remove Int/Uint/Float64Index from tests/arithmetic/test_numeric.py

Author: Terji Petersen

pandas/tests/arithmetic/test_numeric.py

@@ -22,11 +22,7 @@
 )
 import pandas._testing as tm
 from pandas.core import ops
-from pandas.core.api import (
-    Float64Index,
-    Int64Index,
-    UInt64Index,
-)
+from pandas.core.api import NumericIndex
 from pandas.core.computation import expressions as expr
 from pandas.tests.arithmetic.common import (
     assert_invalid_addsub_type,
@@ -375,7 +371,7 @@ def test_divmod_zero(self, zero, numeric_idx):
     @pytest.mark.parametrize("op", [operator.truediv, operator.floordiv])
     def test_div_negative_zero(self, zero, numeric_idx, op):
         # Check that -1 / -0.0 returns np.inf, not -np.inf
-        if isinstance(numeric_idx, UInt64Index):
+        if tm.is_unsigned_integer_dtype(numeric_idx):
             return
         idx = numeric_idx - 3
 
@@ -667,15 +663,15 @@ def test_mul_int_array(self, numeric_idx):
         result = idx * np.array(5, dtype="int64")
         tm.assert_index_equal(result, idx * 5)
 
-        arr_dtype = "uint64" if isinstance(idx, UInt64Index) else "int64"
+        arr_dtype = "uint64" if tm.is_unsigned_integer_dtype(idx) else "int64"
         result = idx * np.arange(5, dtype=arr_dtype)
         tm.assert_index_equal(result, didx)
 
     def test_mul_int_series(self, numeric_idx):
         idx = numeric_idx
         didx = idx * idx
 
-        arr_dtype = "uint64" if isinstance(idx, UInt64Index) else "int64"
+        arr_dtype = "uint64" if tm.is_unsigned_integer_dtype(idx) else "int64"
         result = idx * Series(np.arange(5, dtype=arr_dtype))
         tm.assert_series_equal(result, Series(didx))
 
@@ -712,7 +708,7 @@ def test_pow_float(self, op, numeric_idx, box_with_array):
         # test power calculations both ways, GH#14973
         box = box_with_array
         idx = numeric_idx
-        expected = Float64Index(op(idx.values, 2.0))
+        expected = NumericIndex(op(idx.values, 2.0), dtype=np.float64)
 
         idx = tm.box_expected(idx, box)
         expected = tm.box_expected(expected, box)
@@ -1056,74 +1052,92 @@ def test_series_divmod_zero(self):
 
 
 class TestUFuncCompat:
-    @pytest.mark.parametrize(
-        "holder",
-        [Int64Index, UInt64Index, Float64Index, RangeIndex, Series],
-    )
-    def test_ufunc_compat(self, holder):
-        box = Series if holder is Series else Index
+    @pytest.mark.parametrize("holder", [NumericIndex, RangeIndex, Series])
+    def test_ufunc_compat(self, holder, any_real_numpy_dtype):
+        dtype = any_real_numpy_dtype
 
+        if holder is RangeIndex and dtype != "int64":
+            pytest.skip("Not relevant for RangeIndex")
+
+        box = Series if holder is Series else NumericIndex
         if holder is RangeIndex:
             idx = RangeIndex(0, 5, name="foo")
         else:
-            idx = holder(np.arange(5, dtype="int64"), name="foo")
+            idx = holder(np.arange(5, dtype="int64"), dtype=dtype, name="foo")
+
         result = np.sin(idx)
-        expected = box(np.sin(np.arange(5, dtype="int64")), name="foo")
+
+        if dtype in ["int8", "uint8"]:
+            # index doesn't support float16
+            expected_dtype = np.float32 if holder is NumericIndex else np.float16
+        elif dtype in ["int16", "uint16", "float32"]:
+            expected_dtype = np.float32
+        else:
+            expected_dtype = np.float64
+        exp_val = np.sin(np.arange(5, dtype=dtype)).astype(expected_dtype)
+        expected = box(exp_val, name="foo")
         tm.assert_equal(result, expected)
 
-    @pytest.mark.parametrize("holder", [Int64Index, UInt64Index, Float64Index, Series])
-    def test_ufunc_coercions(self, holder):
-        idx = holder([1, 2, 3, 4, 5], name="x")
-        box = Series if holder is Series else Index
+    @pytest.mark.parametrize("holder", [NumericIndex, Series])
+    def test_ufunc_coercions(self, holder, any_real_numpy_dtype):
+        dtype = any_real_numpy_dtype
+        idx = holder([1, 2, 3, 4, 5], dtype=dtype, name="x")
+        box = Series if holder is Series else NumericIndex
+
+        if dtype in ["int8", "uint8"]:
+            # index doesn't support float16
+            expected_dtype = np.float32 if holder is NumericIndex else np.float16
+        elif dtype in ["int16", "uint16", "float32"]:
+            expected_dtype = np.float32
+        else:
+            expected_dtype = np.float64
 
         result = np.sqrt(idx)
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index(np.sqrt(np.array([1, 2, 3, 4, 5])), name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp_val = np.sqrt(np.array([1, 2, 3, 4, 5], dtype=dtype)).astype(expected_dtype)
+        exp = box(exp_val, name="x")
         tm.assert_equal(result, exp)
 
+        expected_dtype = dtype if tm.is_float_dtype(dtype) else np.float64
+
         result = np.divide(idx, 2.0)
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index([0.5, 1.0, 1.5, 2.0, 2.5], name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp = box([0.5, 1.0, 1.5, 2.0, 2.5], dtype=expected_dtype, name="x")
         tm.assert_equal(result, exp)
 
         # _evaluate_numeric_binop
         result = idx + 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index([3.0, 4.0, 5.0, 6.0, 7.0], name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp = box([3.0, 4.0, 5.0, 6.0, 7.0], dtype=expected_dtype, name="x")
         tm.assert_equal(result, exp)
 
         result = idx - 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index([-1.0, 0.0, 1.0, 2.0, 3.0], name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp = box([-1.0, 0.0, 1.0, 2.0, 3.0], dtype=expected_dtype, name="x")
         tm.assert_equal(result, exp)
 
         result = idx * 1.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index([1.0, 2.0, 3.0, 4.0, 5.0], name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp = box([1.0, 2.0, 3.0, 4.0, 5.0], dtype=expected_dtype, name="x")
         tm.assert_equal(result, exp)
 
         result = idx / 2.0
-        assert result.dtype == "f8" and isinstance(result, box)
-        exp = Float64Index([0.5, 1.0, 1.5, 2.0, 2.5], name="x")
-        exp = tm.box_expected(exp, box)
+        assert result.dtype == expected_dtype and isinstance(result, box)
+        exp = box([0.5, 1.0, 1.5, 2.0, 2.5], dtype=expected_dtype, name="x")
         tm.assert_equal(result, exp)
 
-    @pytest.mark.parametrize("holder", [Int64Index, UInt64Index, Float64Index, Series])
-    def test_ufunc_multiple_return_values(self, holder):
-        obj = holder([1, 2, 3], name="x")
-        box = Series if holder is Series else Index
+    @pytest.mark.parametrize("holder", [NumericIndex, Series])
+    @pytest.mark.parametrize("dtype", ["int64", "uint64", "float64"])
+    def test_ufunc_multiple_return_values(self, holder, dtype):
+        obj = holder([1, 2, 3], dtype=dtype, name="x")
+        box = Series if holder is Series else NumericIndex
 
         result = np.modf(obj)
         assert isinstance(result, tuple)
-        exp1 = Float64Index([0.0, 0.0, 0.0], name="x")
-        exp2 = Float64Index([1.0, 2.0, 3.0], name="x")
-        tm.assert_equal(result[0], tm.box_expected(exp1, box))
-        tm.assert_equal(result[1], tm.box_expected(exp2, box))
+        exp1 = box([0.0, 0.0, 0.0], dtype=np.float64, name="x")
+        exp2 = box([1.0, 2.0, 3.0], dtype=np.float64, name="x")
+        tm.assert_equal(result[0], exp1)
+        tm.assert_equal(result[1], exp2)
 
     def test_ufunc_at(self):
         s = Series([0, 1, 2], index=[1, 2, 3], name="x")
@@ -1214,7 +1228,7 @@ def test_binops_index(self, op, idx1, idx2):
         idx1 = idx1._rename("foo")
         idx2 = idx2._rename("bar")
         result = op(idx1, idx2)
-        expected = op(Int64Index(idx1), Int64Index(idx2))
+        expected = op(NumericIndex(idx1.to_numpy()), NumericIndex(idx2.to_numpy()))
         tm.assert_index_equal(result, expected, exact="equiv")
 
     @pytest.mark.parametrize(
@@ -1239,7 +1253,7 @@ def test_binops_index(self, op, idx1, idx2):
     @pytest.mark.parametrize("scalar", [-1, 1, 2])
     def test_binops_index_scalar(self, op, idx, scalar):
         result = op(idx, scalar)
-        expected = op(Int64Index(idx), scalar)
+        expected = op(NumericIndex(idx.to_numpy(), dtype=np.int64), scalar)
         tm.assert_index_equal(result, expected, exact="equiv")
 
     @pytest.mark.parametrize("idx1", [RangeIndex(0, 10, 1), RangeIndex(0, 20, 2)])
@@ -1250,7 +1264,10 @@ def test_binops_index_pow(self, idx1, idx2):
         idx1 = idx1._rename("foo")
         idx2 = idx2._rename("bar")
         result = pow(idx1, idx2)
-        expected = pow(Int64Index(idx1), Int64Index(idx2))
+        expected = pow(
+            NumericIndex(idx1.to_numpy(), dtype=np.int64),
+            NumericIndex(idx2.to_numpy(), dtype=np.int64),
+        )
         tm.assert_index_equal(result, expected, exact="equiv")
 
     @pytest.mark.parametrize("idx", [RangeIndex(0, 10, 1), RangeIndex(0, 20, 2)])
@@ -1259,7 +1276,7 @@ def test_binops_index_scalar_pow(self, idx, scalar):
         # numpy does not allow powers of negative integers so test separately
         # https://github.com/numpy/numpy/pull/8127
         result = pow(idx, scalar)
-        expected = pow(Int64Index(idx), scalar)
+        expected = pow(NumericIndex(idx.to_numpy(), dtype=np.int64), scalar)
         tm.assert_index_equal(result, expected, exact="equiv")
 
     # TODO: divmod?
@@ -1328,7 +1345,7 @@ def test_numeric_compat2(self):
         # __pow__
         idx = RangeIndex(0, 1000, 2)
         result = idx**2
-        expected = Int64Index(idx._values) ** 2
+        expected = NumericIndex(idx._values, dtype=np.int64) ** 2
         tm.assert_index_equal(Index(result.values), expected, exact=True)
 
     @pytest.mark.parametrize(
@@ -1339,12 +1356,12 @@ def test_numeric_compat2(self):
             (
                 RangeIndex(0, 1000, 1),
                 2,
-                Int64Index(RangeIndex(0, 1000, 1)._values) // 2,
+                NumericIndex(RangeIndex(0, 1000, 1)._values) // 2,
             ),
             (
                 RangeIndex(0, 100, 1),
                 2.0,
-                Int64Index(RangeIndex(0, 100, 1)._values) // 2.0,
+                NumericIndex(RangeIndex(0, 100, 1)._values) // 2.0,
             ),
             (RangeIndex(0), 50, RangeIndex(0)),
             (RangeIndex(2, 4, 2), 3, RangeIndex(0, 1, 1)),