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

Author: Terji Petersen

pandas/core/ops/missing.py

@@ -128,7 +128,8 @@ def mask_zero_div_zero(x, y, result: np.ndarray) -> np.ndarray:
 
         if nan_mask.any() or neginf_mask.any() or posinf_mask.any():
             # Fill negative/0 with -inf, positive/0 with +inf, 0/0 with NaN
-            result = result.astype("float64", copy=False)
+            if not is_float_dtype(result):
+                result = result.astype("float64", copy=False)
 
             result[nan_mask] = np.nan
             result[posinf_mask] = np.inf

pandas/tests/arithmetic/conftest.py

@@ -4,11 +4,7 @@
 import pandas as pd
 from pandas import RangeIndex
 import pandas._testing as tm
-from pandas.core.api import (
-    Float64Index,
-    Int64Index,
-    UInt64Index,
-)
+from pandas.core.api import NumericIndex
 from pandas.core.computation import expressions as expr
 
 
@@ -79,23 +75,23 @@ def zero(request):
     --------
     arr = RangeIndex(5)
     arr / zeros
-    Float64Index([nan, inf, inf, inf, inf], dtype='float64')
+    NumericIndex([nan, inf, inf, inf, inf], dtype='float64')
     """
     return request.param
 
 
 # ------------------------------------------------------------------
 # Vector Fixtures
 
+_numeric_idx_params = [
+    NumericIndex(np.arange(5, dtype=dtype))
+    for dtype in [np.int64, np.uint64, np.float64, np.int8, np.float32]
+] + [RangeIndex(5)]
+
 
 @pytest.fixture(
-    params=[
-        Float64Index(np.arange(5, dtype="float64")),
-        Int64Index(np.arange(5, dtype="int64")),
-        UInt64Index(np.arange(5, dtype="uint64")),
-        RangeIndex(5),
-    ],
-    ids=lambda x: type(x).__name__,
+    params=_numeric_idx_params,
+    ids=lambda x: f"{type(x).__name__}[{x.dtype}]",
 )
 def numeric_idx(request):
     """

pandas/tests/arithmetic/test_numeric.py

@@ -322,51 +322,67 @@ def test_add_sub_datetimedeltalike_invalid(
 
 
 class TestDivisionByZero:
+    def _get_expected_dtype(self, idx, zero):
+        """Helper function for getting the correct return dtype.
+
+        The return dtype for floats depends on the if zero is shaped or not.
+        """
+        if not tm.is_float_dtype(idx) or (hasattr(zero, "shape") and zero.shape):
+            expected_dtype = np.float64
+        else:
+            expected_dtype = idx.dtype
+        return expected_dtype
+
     def test_div_zero(self, zero, numeric_idx):
         idx = numeric_idx
 
-        expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
+        expected_dtype = self._get_expected_dtype(idx, zero=zero)
+        expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=expected_dtype)
         # We only adjust for Index, because Series does not yet apply
         #  the adjustment correctly.
         expected2 = adjust_negative_zero(zero, expected)
 
         result = idx / zero
         tm.assert_index_equal(result, expected2)
         ser_compat = Series(idx).astype("i8") / np.array(zero).astype("i8")
-        tm.assert_series_equal(ser_compat, Series(expected))
+        tm.assert_series_equal(ser_compat, Series(expected.astype(np.float64)))
 
     def test_floordiv_zero(self, zero, numeric_idx):
         idx = numeric_idx
 
-        expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
+        expected_dtype = self._get_expected_dtype(idx, zero=zero)
+        expected = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=expected_dtype)
         # We only adjust for Index, because Series does not yet apply
         #  the adjustment correctly.
         expected2 = adjust_negative_zero(zero, expected)
 
         result = idx // zero
         tm.assert_index_equal(result, expected2)
         ser_compat = Series(idx).astype("i8") // np.array(zero).astype("i8")
-        tm.assert_series_equal(ser_compat, Series(expected))
+        tm.assert_series_equal(ser_compat, Series(expected.astype(np.float64)))
 
     def test_mod_zero(self, zero, numeric_idx):
         idx = numeric_idx
 
-        expected = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=np.float64)
+        expected_dtype = self._get_expected_dtype(idx, zero=zero)
+        expected = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=expected_dtype)
+
         result = idx % zero
         tm.assert_index_equal(result, expected)
         ser_compat = Series(idx).astype("i8") % np.array(zero).astype("i8")
-        tm.assert_series_equal(ser_compat, Series(result))
+        tm.assert_series_equal(ser_compat, Series(result, dtype=np.float64))
 
     def test_divmod_zero(self, zero, numeric_idx):
         idx = numeric_idx
 
-        exleft = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=np.float64)
-        exright = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=np.float64)
-        exleft = adjust_negative_zero(zero, exleft)
+        expe_dtype = self._get_expected_dtype(idx, zero=zero)
+        exp_left = Index([np.nan, np.inf, np.inf, np.inf, np.inf], dtype=expe_dtype)
+        exp_right = Index([np.nan, np.nan, np.nan, np.nan, np.nan], dtype=expe_dtype)
+        exp_left = adjust_negative_zero(zero, exp_left)
 
         result = divmod(idx, zero)
-        tm.assert_index_equal(result[0], exleft)
-        tm.assert_index_equal(result[1], exright)
+        tm.assert_index_equal(result[0], exp_left)
+        tm.assert_index_equal(result[1], exp_right)
 
     @pytest.mark.parametrize("op", [operator.truediv, operator.floordiv])
     def test_div_negative_zero(self, zero, numeric_idx, op):
@@ -375,7 +391,8 @@ def test_div_negative_zero(self, zero, numeric_idx, op):
             return
         idx = numeric_idx - 3
 
-        expected = Index([-np.inf, -np.inf, -np.inf, np.nan, np.inf], dtype=np.float64)
+        exp_dtype = self._get_expected_dtype(idx, zero=zero)
+        expected = Index([-np.inf, -np.inf, -np.inf, np.nan, np.inf], dtype=exp_dtype)
         expected = adjust_negative_zero(zero, expected)
 
         result = op(idx, zero)
@@ -641,7 +658,8 @@ def test_div_equiv_binop(self):
     def test_div_int(self, numeric_idx):
         idx = numeric_idx
         result = idx / 1
-        expected = idx.astype("float64")
+        expected_dtype = idx.dtype if tm.is_float_dtype(idx) else np.float64
+        expected = idx.astype(expected_dtype)
         tm.assert_index_equal(result, expected)
 
         result = idx / 2
@@ -663,16 +681,14 @@ 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 tm.is_unsigned_integer_dtype(idx) else "int64"
-        result = idx * np.arange(5, dtype=arr_dtype)
+        result = idx * np.arange(5, dtype=idx.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 tm.is_unsigned_integer_dtype(idx) else "int64"
-        result = idx * Series(np.arange(5, dtype=arr_dtype))
+        result = idx * Series(np.arange(5, dtype=idx.dtype))
         tm.assert_series_equal(result, Series(didx))
 
     def test_mul_float_series(self, numeric_idx):
@@ -708,7 +724,8 @@ 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 = NumericIndex(op(idx.values, 2.0), dtype=np.float64)
+        expected_dtype = idx.dtype if tm.is_float_dtype(idx) else np.float64
+        expected = NumericIndex(op(idx.values, 2.0), dtype=expected_dtype)
 
         idx = tm.box_expected(idx, box)
         expected = tm.box_expected(expected, box)