import numpy as np import pytest from pandas.core.dtypes.generic import ABCIndexClass import pandas as pd from pandas.api.types import is_float, is_float_dtype, is_integer, is_scalar from pandas.core.arrays import IntegerArray, integer_array from pandas.core.arrays.integer import ( Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, UInt8Dtype, UInt16Dtype, UInt32Dtype, UInt64Dtype, ) from pandas.tests.extension.base import BaseOpsUtil import pandas.util.testing as tm def make_data(): return list(range(8)) + [np.nan] + list(range(10, 98)) + [np.nan] + [99, 100] @pytest.fixture( params=[ Int8Dtype, Int16Dtype, Int32Dtype, Int64Dtype, UInt8Dtype, UInt16Dtype, UInt32Dtype, UInt64Dtype, ] ) def dtype(request): return request.param() @pytest.fixture def data(dtype): return integer_array(make_data(), dtype=dtype) @pytest.fixture def data_missing(dtype): return integer_array([np.nan, 1], dtype=dtype) @pytest.fixture(params=["data", "data_missing"]) def all_data(request, data, data_missing): """Parametrized fixture giving 'data' and 'data_missing'""" if request.param == "data": return data elif request.param == "data_missing": return data_missing def test_dtypes(dtype): # smoke tests on auto dtype construction if dtype.is_signed_integer: assert np.dtype(dtype.type).kind == "i" else: assert np.dtype(dtype.type).kind == "u" assert dtype.name is not None @pytest.mark.parametrize( "dtype, expected", [ (Int8Dtype(), "Int8Dtype()"), (Int16Dtype(), "Int16Dtype()"), (Int32Dtype(), "Int32Dtype()"), (Int64Dtype(), "Int64Dtype()"), (UInt8Dtype(), "UInt8Dtype()"), (UInt16Dtype(), "UInt16Dtype()"), (UInt32Dtype(), "UInt32Dtype()"), (UInt64Dtype(), "UInt64Dtype()"), ], ) def test_repr_dtype(dtype, expected): assert repr(dtype) == expected def test_repr_array(): result = repr(integer_array([1, None, 3])) expected = "\n[1, NaN, 3]\nLength: 3, dtype: Int64" assert result == expected def test_repr_array_long(): data = integer_array([1, 2, None] * 1000) expected = ( "\n" "[ 1, 2, NaN, 1, 2, NaN, 1, 2, NaN, 1,\n" " ...\n" " NaN, 1, 2, NaN, 1, 2, NaN, 1, 2, NaN]\n" "Length: 3000, dtype: Int64" ) result = repr(data) assert result == expected class TestConstructors: def test_from_dtype_from_float(self, data): # construct from our dtype & string dtype dtype = data.dtype # from float expected = pd.Series(data) result = pd.Series(np.array(data).astype("float"), dtype=str(dtype)) tm.assert_series_equal(result, expected) # from int / list expected = pd.Series(data) result = pd.Series(np.array(data).tolist(), dtype=str(dtype)) tm.assert_series_equal(result, expected) # from int / array expected = pd.Series(data).dropna().reset_index(drop=True) dropped = np.array(data.dropna()).astype(np.dtype((dtype.type))) result = pd.Series(dropped, dtype=str(dtype)) tm.assert_series_equal(result, expected) class TestArithmeticOps(BaseOpsUtil): def _check_divmod_op(self, s, op, other, exc=None): super()._check_divmod_op(s, op, other, None) def _check_op(self, s, op_name, other, exc=None): op = self.get_op_from_name(op_name) result = op(s, other) # compute expected mask = s.isna() # if s is a DataFrame, squeeze to a Series # for comparison if isinstance(s, pd.DataFrame): result = result.squeeze() s = s.squeeze() mask = mask.squeeze() # other array is an Integer if isinstance(other, IntegerArray): omask = getattr(other, "mask", None) mask = getattr(other, "data", other) if omask is not None: mask |= omask # 1 ** na is na, so need to unmask those if op_name == "__pow__": mask = np.where(s == 1, False, mask) elif op_name == "__rpow__": mask = np.where(other == 1, False, mask) # float result type or float op if ( is_float_dtype(other) or is_float(other) or op_name in ["__rtruediv__", "__truediv__", "__rdiv__", "__div__"] ): rs = s.astype("float") expected = op(rs, other) self._check_op_float(result, expected, mask, s, op_name, other) # integer result type else: rs = pd.Series(s.values._data, name=s.name) expected = op(rs, other) self._check_op_integer(result, expected, mask, s, op_name, other) def _check_op_float(self, result, expected, mask, s, op_name, other): # check comparisons that are resulting in float dtypes expected[mask] = np.nan if "floordiv" in op_name: # Series op sets 1//0 to np.inf, which IntegerArray does not do (yet) mask2 = np.isinf(expected) & np.isnan(result) expected[mask2] = np.nan tm.assert_series_equal(result, expected) def _check_op_integer(self, result, expected, mask, s, op_name, other): # check comparisons that are resulting in integer dtypes # to compare properly, we convert the expected # to float, mask to nans and convert infs # if we have uints then we process as uints # then conert to float # and we ultimately want to create a IntArray # for comparisons fill_value = 0 # mod/rmod turn floating 0 into NaN while # integer works as expected (no nan) if op_name in ["__mod__", "__rmod__"]: if is_scalar(other): if other == 0: expected[s.values == 0] = 0 else: expected = expected.fillna(0) else: expected[(s.values == 0) & ((expected == 0) | expected.isna())] = 0 try: expected[(expected == np.inf) | (expected == -np.inf)] = fill_value original = expected expected = expected.astype(s.dtype) except ValueError: expected = expected.astype(float) expected[(expected == np.inf) | (expected == -np.inf)] = fill_value original = expected expected = expected.astype(s.dtype) expected[mask] = np.nan # assert that the expected astype is ok # (skip for unsigned as they have wrap around) if not s.dtype.is_unsigned_integer: original = pd.Series(original) # we need to fill with 0's to emulate what an astype('int') does # (truncation) for certain ops if op_name in ["__rtruediv__", "__rdiv__"]: mask |= original.isna() original = original.fillna(0).astype("int") original = original.astype("float") original[mask] = np.nan tm.assert_series_equal(original, expected.astype("float")) # assert our expected result tm.assert_series_equal(result, expected) def test_arith_integer_array(self, data, all_arithmetic_operators): # we operate with a rhs of an integer array op = all_arithmetic_operators s = pd.Series(data) rhs = pd.Series([1] * len(data), dtype=data.dtype) rhs.iloc[-1] = np.nan self._check_op(s, op, rhs) def test_arith_series_with_scalar(self, data, all_arithmetic_operators): # scalar op = all_arithmetic_operators s = pd.Series(data) self._check_op(s, op, 1, exc=TypeError) def test_arith_frame_with_scalar(self, data, all_arithmetic_operators): # frame & scalar op = all_arithmetic_operators df = pd.DataFrame({"A": data}) self._check_op(df, op, 1, exc=TypeError) def test_arith_series_with_array(self, data, all_arithmetic_operators): # ndarray & other series op = all_arithmetic_operators s = pd.Series(data) other = np.ones(len(s), dtype=s.dtype.type) self._check_op(s, op, other, exc=TypeError) def test_arith_coerce_scalar(self, data, all_arithmetic_operators): op = all_arithmetic_operators s = pd.Series(data) other = 0.01 self._check_op(s, op, other) @pytest.mark.parametrize("other", [1.0, 1.0, np.array(1.0), np.array([1.0])]) def test_arithmetic_conversion(self, all_arithmetic_operators, other): # if we have a float operand we should have a float result # if that is equal to an integer op = self.get_op_from_name(all_arithmetic_operators) s = pd.Series([1, 2, 3], dtype="Int64") result = op(s, other) assert result.dtype is np.dtype("float") @pytest.mark.parametrize("other", [0, 0.5]) def test_arith_zero_dim_ndarray(self, other): arr = integer_array([1, None, 2]) result = arr + np.array(other) expected = arr + other tm.assert_equal(result, expected) def test_error(self, data, all_arithmetic_operators): # invalid ops op = all_arithmetic_operators s = pd.Series(data) ops = getattr(s, op) opa = getattr(data, op) # invalid scalars with pytest.raises(TypeError): ops("foo") with pytest.raises(TypeError): ops(pd.Timestamp("20180101")) # invalid array-likes with pytest.raises(TypeError): ops(pd.Series("foo", index=s.index)) if op != "__rpow__": # TODO(extension) # rpow with a datetimelike coerces the integer array incorrectly with pytest.raises(TypeError): ops(pd.Series(pd.date_range("20180101", periods=len(s)))) # 2d with pytest.raises(NotImplementedError): opa(pd.DataFrame({"A": s})) with pytest.raises(NotImplementedError): opa(np.arange(len(s)).reshape(-1, len(s))) def test_pow(self): # https://github.com/pandas-dev/pandas/issues/22022 a = integer_array([1, np.nan, np.nan, 1]) b = integer_array([1, np.nan, 1, np.nan]) result = a ** b expected = pd.core.arrays.integer_array([1, np.nan, np.nan, 1]) tm.assert_extension_array_equal(result, expected) def test_rpow_one_to_na(self): # https://github.com/pandas-dev/pandas/issues/22022 arr = integer_array([np.nan, np.nan]) result = np.array([1.0, 2.0]) ** arr expected = np.array([1.0, np.nan]) tm.assert_numpy_array_equal(result, expected) 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)) # fill the nan locations expected[data._mask] = op_name == "__ne__" tm.assert_series_equal(result, expected) # series s = pd.Series(data) result = op(s, other) expected = pd.Series(data._data) expected = op(expected, other) # fill the nan locations expected[data._mask] = op_name == "__ne__" tm.assert_series_equal(result, expected) def test_compare_scalar(self, data, all_compare_operators): op_name = all_compare_operators self._compare_other(data, op_name, 0) def test_compare_array(self, data, all_compare_operators): op_name = all_compare_operators other = pd.Series([0] * len(data)) self._compare_other(data, op_name, other) class TestCasting: pass @pytest.mark.parametrize("dropna", [True, False]) def test_construct_index(self, all_data, dropna): # ensure that we do not coerce to Float64Index, rather # keep as Index all_data = all_data[:10] if dropna: other = np.array(all_data[~all_data.isna()]) else: other = all_data result = pd.Index(integer_array(other, dtype=all_data.dtype)) expected = pd.Index(other, dtype=object) tm.assert_index_equal(result, expected) @pytest.mark.parametrize("dropna", [True, False]) def test_astype_index(self, all_data, dropna): # as an int/uint index to Index all_data = all_data[:10] if dropna: other = all_data[~all_data.isna()] else: other = all_data dtype = all_data.dtype idx = pd.Index(np.array(other)) assert isinstance(idx, ABCIndexClass) result = idx.astype(dtype) expected = idx.astype(object).astype(dtype) tm.assert_index_equal(result, expected) def test_astype(self, all_data): all_data = all_data[:10] ints = all_data[~all_data.isna()] mixed = all_data dtype = Int8Dtype() # coerce to same type - ints s = pd.Series(ints) result = s.astype(all_data.dtype) expected = pd.Series(ints) tm.assert_series_equal(result, expected) # coerce to same other - ints s = pd.Series(ints) result = s.astype(dtype) expected = pd.Series(ints, dtype=dtype) tm.assert_series_equal(result, expected) # coerce to same numpy_dtype - ints s = pd.Series(ints) result = s.astype(all_data.dtype.numpy_dtype) expected = pd.Series(ints._data.astype(all_data.dtype.numpy_dtype)) tm.assert_series_equal(result, expected) # coerce to same type - mixed s = pd.Series(mixed) result = s.astype(all_data.dtype) expected = pd.Series(mixed) tm.assert_series_equal(result, expected) # coerce to same other - mixed s = pd.Series(mixed) result = s.astype(dtype) expected = pd.Series(mixed, dtype=dtype) tm.assert_series_equal(result, expected) # coerce to same numpy_dtype - mixed s = pd.Series(mixed) with pytest.raises(ValueError): s.astype(all_data.dtype.numpy_dtype) # coerce to object s = pd.Series(mixed) result = s.astype("object") expected = pd.Series(np.asarray(mixed)) tm.assert_series_equal(result, expected) @pytest.mark.parametrize("dtype", [Int8Dtype(), "Int8", UInt32Dtype(), "UInt32"]) def test_astype_specific_casting(self, dtype): s = pd.Series([1, 2, 3], dtype="Int64") result = s.astype(dtype) expected = pd.Series([1, 2, 3], dtype=dtype) tm.assert_series_equal(result, expected) s = pd.Series([1, 2, 3, None], dtype="Int64") result = s.astype(dtype) expected = pd.Series([1, 2, 3, None], dtype=dtype) tm.assert_series_equal(result, expected) def test_construct_cast_invalid(self, dtype): msg = "cannot safely" arr = [1.2, 2.3, 3.7] with pytest.raises(TypeError, match=msg): integer_array(arr, dtype=dtype) with pytest.raises(TypeError, match=msg): pd.Series(arr).astype(dtype) arr = [1.2, 2.3, 3.7, np.nan] with pytest.raises(TypeError, match=msg): integer_array(arr, dtype=dtype) with pytest.raises(TypeError, match=msg): pd.Series(arr).astype(dtype) def test_frame_repr(data_missing): df = pd.DataFrame({"A": data_missing}) result = repr(df) expected = " A\n0 NaN\n1 1" assert result == expected def test_conversions(data_missing): # astype to object series df = pd.DataFrame({"A": data_missing}) result = df["A"].astype("object") expected = pd.Series(np.array([np.nan, 1], dtype=object), name="A") tm.assert_series_equal(result, expected) # convert to object ndarray # we assert that we are exactly equal # including type conversions of scalars result = df["A"].astype("object").values expected = np.array([np.nan, 1], dtype=object) tm.assert_numpy_array_equal(result, expected) for r, e in zip(result, expected): if pd.isnull(r): assert pd.isnull(e) elif is_integer(r): assert r == e assert is_integer(e) else: assert r == e assert type(r) == type(e) def test_integer_array_constructor(): values = np.array([1, 2, 3, 4], dtype="int64") mask = np.array([False, False, False, True], dtype="bool") result = IntegerArray(values, mask) expected = integer_array([1, 2, 3, np.nan], dtype="int64") tm.assert_extension_array_equal(result, expected) with pytest.raises(TypeError): IntegerArray(values.tolist(), mask) with pytest.raises(TypeError): IntegerArray(values, mask.tolist()) with pytest.raises(TypeError): IntegerArray(values.astype(float), mask) with pytest.raises(TypeError): IntegerArray(values) @pytest.mark.parametrize( "a, b", [ ([1, None], [1, np.nan]), ([None], [np.nan]), ([None, np.nan], [np.nan, np.nan]), ([np.nan, np.nan], [np.nan, np.nan]), ], ) def test_integer_array_constructor_none_is_nan(a, b): result = integer_array(a) expected = integer_array(b) tm.assert_extension_array_equal(result, expected) def test_integer_array_constructor_copy(): values = np.array([1, 2, 3, 4], dtype="int64") mask = np.array([False, False, False, True], dtype="bool") result = IntegerArray(values, mask) assert result._data is values assert result._mask is mask result = IntegerArray(values, mask, copy=True) assert result._data is not values assert result._mask is not mask @pytest.mark.parametrize( "values", [ ["foo", "bar"], ["1", "2"], "foo", 1, 1.0, pd.date_range("20130101", periods=2), np.array(["foo"]), [[1, 2], [3, 4]], [np.nan, {"a": 1}], ], ) def test_to_integer_array_error(values): # error in converting existing arrays to IntegerArrays with pytest.raises(TypeError): integer_array(values) def test_to_integer_array_inferred_dtype(): # if values has dtype -> respect it result = integer_array(np.array([1, 2], dtype="int8")) assert result.dtype == Int8Dtype() result = integer_array(np.array([1, 2], dtype="int32")) assert result.dtype == Int32Dtype() # if values have no dtype -> always int64 result = integer_array([1, 2]) assert result.dtype == Int64Dtype() def test_to_integer_array_dtype_keyword(): result = integer_array([1, 2], dtype="int8") assert result.dtype == Int8Dtype() # if values has dtype -> override it result = integer_array(np.array([1, 2], dtype="int8"), dtype="int32") assert result.dtype == Int32Dtype() def test_to_integer_array_float(): result = integer_array([1.0, 2.0]) expected = integer_array([1, 2]) tm.assert_extension_array_equal(result, expected) with pytest.raises(TypeError, match="cannot safely cast non-equivalent"): integer_array([1.5, 2.0]) # for float dtypes, the itemsize is not preserved result = integer_array(np.array([1.0, 2.0], dtype="float32")) assert result.dtype == Int64Dtype() @pytest.mark.parametrize( "bool_values, int_values, target_dtype, expected_dtype", [ ([False, True], [0, 1], Int64Dtype(), Int64Dtype()), ([False, True], [0, 1], "Int64", Int64Dtype()), ([False, True, np.nan], [0, 1, np.nan], Int64Dtype(), Int64Dtype()), ], ) def test_to_integer_array_bool(bool_values, int_values, target_dtype, expected_dtype): result = integer_array(bool_values, dtype=target_dtype) assert result.dtype == expected_dtype expected = integer_array(int_values, dtype=target_dtype) tm.assert_extension_array_equal(result, expected) @pytest.mark.parametrize( "values, to_dtype, result_dtype", [ (np.array([1], dtype="int64"), None, Int64Dtype), (np.array([1, np.nan]), None, Int64Dtype), (np.array([1, np.nan]), "int8", Int8Dtype), ], ) def test_to_integer_array(values, to_dtype, result_dtype): # convert existing arrays to IntegerArrays result = integer_array(values, dtype=to_dtype) assert result.dtype == result_dtype() expected = integer_array(values, dtype=result_dtype()) tm.assert_extension_array_equal(result, expected) def test_cross_type_arithmetic(): df = pd.DataFrame( { "A": pd.Series([1, 2, np.nan], dtype="Int64"), "B": pd.Series([1, np.nan, 3], dtype="UInt8"), "C": [1, 2, 3], } ) result = df.A + df.C expected = pd.Series([2, 4, np.nan], dtype="Int64") tm.assert_series_equal(result, expected) result = (df.A + df.C) * 3 == 12 expected = pd.Series([False, True, False]) tm.assert_series_equal(result, expected) result = df.A + df.B expected = pd.Series([2, np.nan, np.nan], dtype="Int64") tm.assert_series_equal(result, expected) @pytest.mark.parametrize("op", ["sum", "min", "max", "prod"]) def test_preserve_dtypes(op): # TODO(#22346): preserve Int64 dtype # for ops that enable (mean would actually work here # but generally it is a float return value) df = pd.DataFrame( { "A": ["a", "b", "b"], "B": [1, None, 3], "C": integer_array([1, None, 3], dtype="Int64"), } ) # op result = getattr(df.C, op)() assert isinstance(result, int) # groupby result = getattr(df.groupby("A"), op)() expected = pd.DataFrame( {"B": np.array([1.0, 3.0]), "C": integer_array([1, 3], dtype="Int64")}, index=pd.Index(["a", "b"], name="A"), ) tm.assert_frame_equal(result, expected) @pytest.mark.parametrize("op", ["mean"]) def test_reduce_to_float(op): # some reduce ops always return float, even if the result # is a rounded number df = pd.DataFrame( { "A": ["a", "b", "b"], "B": [1, None, 3], "C": integer_array([1, None, 3], dtype="Int64"), } ) # op result = getattr(df.C, op)() assert isinstance(result, float) # groupby result = getattr(df.groupby("A"), op)() expected = pd.DataFrame( {"B": np.array([1.0, 3.0]), "C": integer_array([1, 3], dtype="Int64")}, index=pd.Index(["a", "b"], name="A"), ) tm.assert_frame_equal(result, expected) def test_astype_nansafe(): # see gh-22343 arr = integer_array([np.nan, 1, 2], dtype="Int8") msg = "cannot convert float NaN to integer" with pytest.raises(ValueError, match=msg): arr.astype("uint32") @pytest.mark.parametrize("ufunc", [np.abs, np.sign]) def test_ufuncs_single_int(ufunc): a = integer_array([1, 2, -3, np.nan]) result = ufunc(a) expected = integer_array(ufunc(a.astype(float))) tm.assert_extension_array_equal(result, expected) s = pd.Series(a) result = ufunc(s) expected = pd.Series(integer_array(ufunc(a.astype(float)))) tm.assert_series_equal(result, expected) @pytest.mark.parametrize("ufunc", [np.log, np.exp, np.sin, np.cos, np.sqrt]) def test_ufuncs_single_float(ufunc): a = integer_array([1, 2, -3, np.nan]) with np.errstate(invalid="ignore"): result = ufunc(a) expected = ufunc(a.astype(float)) tm.assert_numpy_array_equal(result, expected) s = pd.Series(a) with np.errstate(invalid="ignore"): result = ufunc(s) expected = ufunc(s.astype(float)) tm.assert_series_equal(result, expected) @pytest.mark.parametrize("ufunc", [np.add, np.subtract]) def test_ufuncs_binary_int(ufunc): # two IntegerArrays a = integer_array([1, 2, -3, np.nan]) result = ufunc(a, a) expected = integer_array(ufunc(a.astype(float), a.astype(float))) tm.assert_extension_array_equal(result, expected) # IntegerArray with numpy array arr = np.array([1, 2, 3, 4]) result = ufunc(a, arr) expected = integer_array(ufunc(a.astype(float), arr)) tm.assert_extension_array_equal(result, expected) result = ufunc(arr, a) expected = integer_array(ufunc(arr, a.astype(float))) tm.assert_extension_array_equal(result, expected) # IntegerArray with scalar result = ufunc(a, 1) expected = integer_array(ufunc(a.astype(float), 1)) tm.assert_extension_array_equal(result, expected) result = ufunc(1, a) expected = integer_array(ufunc(1, a.astype(float))) tm.assert_extension_array_equal(result, expected) @pytest.mark.parametrize("values", [[0, 1], [0, None]]) def test_ufunc_reduce_raises(values): a = integer_array(values) with pytest.raises(NotImplementedError): np.add.reduce(a) # TODO(jreback) - these need testing / are broken # shift # set_index (destroys type)