22
22
)
23
23
import pandas ._testing as tm
24
24
from pandas .core import ops
25
- from pandas .core .api import (
26
- Float64Index ,
27
- Int64Index ,
28
- UInt64Index ,
29
- )
25
+ from pandas .core .api import NumericIndex
30
26
from pandas .core .computation import expressions as expr
31
27
from pandas .tests .arithmetic .common import (
32
28
assert_invalid_addsub_type ,
@@ -1058,72 +1054,77 @@ def test_series_divmod_zero(self):
1058
1054
1059
1055
1060
1056
class TestUFuncCompat :
1061
- @pytest .mark .parametrize (
1062
- "holder" ,
1063
- [Int64Index , UInt64Index , Float64Index , RangeIndex , Series ],
1064
- )
1065
- def test_ufunc_compat (self , holder ):
1057
+ # TODO: add more dtypes
1058
+ @pytest .mark .parametrize ("holder" , [NumericIndex , RangeIndex , Series ])
1059
+ @pytest .mark .parametrize ("dtype" , [np .int64 , np .uint64 , np .float64 ])
1060
+ def test_ufunc_compat (self , holder , dtype ):
1066
1061
box = Series if holder is Series else Index
1067
1062
1068
1063
if holder is RangeIndex :
1064
+ if dtype != np .int64 :
1065
+ pytest .skip (f"dtype { dtype } not relevant for RangeIndex" )
1069
1066
idx = RangeIndex (0 , 5 , name = "foo" )
1070
1067
else :
1071
- idx = holder (np .arange (5 , dtype = "int64" ), name = "foo" )
1068
+ idx = holder (np .arange (5 , dtype = dtype ), name = "foo" )
1072
1069
result = np .sin (idx )
1073
- expected = box (np .sin (np .arange (5 , dtype = "int64" )), name = "foo" )
1070
+ expected = box (np .sin (np .arange (5 , dtype = dtype )), name = "foo" )
1074
1071
tm .assert_equal (result , expected )
1075
1072
1076
- @pytest .mark .parametrize ("holder" , [Int64Index , UInt64Index , Float64Index , Series ])
1077
- def test_ufunc_coercions (self , holder ):
1078
- idx = holder ([1 , 2 , 3 , 4 , 5 ], name = "x" )
1073
+ # TODO: add more dtypes
1074
+ @pytest .mark .parametrize ("holder" , [NumericIndex , Series ])
1075
+ @pytest .mark .parametrize ("dtype" , [np .int64 , np .uint64 , np .float64 ])
1076
+ def test_ufunc_coercions (self , holder , dtype ):
1077
+ idx = holder ([1 , 2 , 3 , 4 , 5 ], dtype = dtype , name = "x" )
1079
1078
box = Series if holder is Series else Index
1080
1079
1081
1080
result = np .sqrt (idx )
1082
1081
assert result .dtype == "f8" and isinstance (result , box )
1083
- exp = Float64Index (np .sqrt (np .array ([1 , 2 , 3 , 4 , 5 ])), name = "x" )
1082
+ exp = Index (np .sqrt (np .array ([1 , 2 , 3 , 4 , 5 ], dtype = np . float64 )), name = "x" )
1084
1083
exp = tm .box_expected (exp , box )
1085
1084
tm .assert_equal (result , exp )
1086
1085
1087
1086
result = np .divide (idx , 2.0 )
1088
1087
assert result .dtype == "f8" and isinstance (result , box )
1089
- exp = Float64Index ([0.5 , 1.0 , 1.5 , 2.0 , 2.5 ], name = "x" )
1088
+ exp = Index ([0.5 , 1.0 , 1.5 , 2.0 , 2.5 ], dtype = np . float64 , name = "x" )
1090
1089
exp = tm .box_expected (exp , box )
1091
1090
tm .assert_equal (result , exp )
1092
1091
1093
1092
# _evaluate_numeric_binop
1094
1093
result = idx + 2.0
1095
1094
assert result .dtype == "f8" and isinstance (result , box )
1096
- exp = Float64Index ([3.0 , 4.0 , 5.0 , 6.0 , 7.0 ], name = "x" )
1095
+ exp = Index ([3.0 , 4.0 , 5.0 , 6.0 , 7.0 ], dtype = np . float64 , name = "x" )
1097
1096
exp = tm .box_expected (exp , box )
1098
1097
tm .assert_equal (result , exp )
1099
1098
1100
1099
result = idx - 2.0
1101
1100
assert result .dtype == "f8" and isinstance (result , box )
1102
- exp = Float64Index ([- 1.0 , 0.0 , 1.0 , 2.0 , 3.0 ], name = "x" )
1101
+ exp = Index ([- 1.0 , 0.0 , 1.0 , 2.0 , 3.0 ], dtype = np . float64 , name = "x" )
1103
1102
exp = tm .box_expected (exp , box )
1104
1103
tm .assert_equal (result , exp )
1105
1104
1106
1105
result = idx * 1.0
1107
1106
assert result .dtype == "f8" and isinstance (result , box )
1108
- exp = Float64Index ([1.0 , 2.0 , 3.0 , 4.0 , 5.0 ], name = "x" )
1107
+ exp = Index ([1.0 , 2.0 , 3.0 , 4.0 , 5.0 ], dtype = np . float64 , name = "x" )
1109
1108
exp = tm .box_expected (exp , box )
1110
1109
tm .assert_equal (result , exp )
1111
1110
1112
1111
result = idx / 2.0
1113
1112
assert result .dtype == "f8" and isinstance (result , box )
1114
- exp = Float64Index ([0.5 , 1.0 , 1.5 , 2.0 , 2.5 ], name = "x" )
1113
+ exp = Index ([0.5 , 1.0 , 1.5 , 2.0 , 2.5 ], dtype = np . float64 , name = "x" )
1115
1114
exp = tm .box_expected (exp , box )
1116
1115
tm .assert_equal (result , exp )
1117
1116
1118
- @pytest .mark .parametrize ("holder" , [Int64Index , UInt64Index , Float64Index , Series ])
1119
- def test_ufunc_multiple_return_values (self , holder ):
1120
- obj = holder ([1 , 2 , 3 ], name = "x" )
1117
+ # TODO: add more dtypes
1118
+ @pytest .mark .parametrize ("holder" , [NumericIndex , Series ])
1119
+ @pytest .mark .parametrize ("dtype" , [np .int64 , np .uint64 , np .float64 ])
1120
+ def test_ufunc_multiple_return_values (self , holder , dtype ):
1121
+ obj = holder ([1 , 2 , 3 ], dtype = dtype , name = "x" )
1121
1122
box = Series if holder is Series else Index
1122
1123
1123
1124
result = np .modf (obj )
1124
1125
assert isinstance (result , tuple )
1125
- exp1 = Float64Index ([0.0 , 0.0 , 0.0 ], name = "x" )
1126
- exp2 = Float64Index ([1.0 , 2.0 , 3.0 ], name = "x" )
1126
+ exp1 = Index ([0.0 , 0.0 , 0.0 ], dtype = np . float64 , name = "x" )
1127
+ exp2 = Index ([1.0 , 2.0 , 3.0 ], dtype = np . float64 , name = "x" )
1127
1128
tm .assert_equal (result [0 ], tm .box_expected (exp1 , box ))
1128
1129
tm .assert_equal (result [1 ], tm .box_expected (exp2 , box ))
1129
1130
@@ -1241,7 +1242,7 @@ def test_binops_index(self, op, idx1, idx2):
1241
1242
@pytest .mark .parametrize ("scalar" , [- 1 , 1 , 2 ])
1242
1243
def test_binops_index_scalar (self , op , idx , scalar ):
1243
1244
result = op (idx , scalar )
1244
- expected = op (Int64Index (idx ), scalar )
1245
+ expected = op (Index (idx . to_numpy () ), scalar )
1245
1246
tm .assert_index_equal (result , expected , exact = "equiv" )
1246
1247
1247
1248
@pytest .mark .parametrize ("idx1" , [RangeIndex (0 , 10 , 1 ), RangeIndex (0 , 20 , 2 )])
@@ -1261,7 +1262,7 @@ def test_binops_index_scalar_pow(self, idx, scalar):
1261
1262
# numpy does not allow powers of negative integers so test separately
1262
1263
# https://github.com/numpy/numpy/pull/8127
1263
1264
result = pow (idx , scalar )
1264
- expected = pow (Int64Index (idx ), scalar )
1265
+ expected = pow (Index (idx . to_numpy () ), scalar )
1265
1266
tm .assert_index_equal (result , expected , exact = "equiv" )
1266
1267
1267
1268
# TODO: divmod?
@@ -1336,17 +1337,18 @@ def test_numeric_compat2(self):
1336
1337
@pytest .mark .parametrize (
1337
1338
"idx, div, expected" ,
1338
1339
[
1340
+ # TODO: add more dtypes
1339
1341
(RangeIndex (0 , 1000 , 2 ), 2 , RangeIndex (0 , 500 , 1 )),
1340
1342
(RangeIndex (- 99 , - 201 , - 3 ), - 3 , RangeIndex (33 , 67 , 1 )),
1341
1343
(
1342
1344
RangeIndex (0 , 1000 , 1 ),
1343
1345
2 ,
1344
- Int64Index (RangeIndex (0 , 1000 , 1 )._values ) // 2 ,
1346
+ Index (RangeIndex (0 , 1000 , 1 )._values ) // 2 ,
1345
1347
),
1346
1348
(
1347
1349
RangeIndex (0 , 100 , 1 ),
1348
1350
2.0 ,
1349
- Int64Index (RangeIndex (0 , 100 , 1 )._values ) // 2.0 ,
1351
+ Index (RangeIndex (0 , 100 , 1 )._values ) // 2.0 ,
1350
1352
),
1351
1353
(RangeIndex (0 ), 50 , RangeIndex (0 )),
1352
1354
(RangeIndex (2 , 4 , 2 ), 3 , RangeIndex (0 , 1 , 1 )),
0 commit comments