diff --git a/doc/source/whatsnew/v1.1.0.rst b/doc/source/whatsnew/v1.1.0.rst index 21e59805fa143..9b0e137e63bd7 100644 --- a/doc/source/whatsnew/v1.1.0.rst +++ b/doc/source/whatsnew/v1.1.0.rst @@ -382,7 +382,7 @@ Sparse ExtensionArray ^^^^^^^^^^^^^^ -- +- integer EA will now return objects of native type only when iterated. (:issue:`29738`) - diff --git a/pandas/core/arrays/integer.py b/pandas/core/arrays/integer.py index fb33840ad757c..adb36e00aac37 100644 --- a/pandas/core/arrays/integer.py +++ b/pandas/core/arrays/integer.py @@ -354,6 +354,13 @@ def __init__(self, values: np.ndarray, mask: np.ndarray, copy: bool = False): ) super().__init__(values, mask, copy=copy) + def __iter__(self): + for i in range(len(self)): + if self._mask[i]: + yield self.dtype.na_value + else: + yield self._data[i].item() + @classmethod def _from_sequence(cls, scalars, dtype=None, copy: bool = False) -> "IntegerArray": return integer_array(scalars, dtype=dtype, copy=copy) diff --git a/pandas/tests/arrays/test_integer.py b/pandas/tests/arrays/test_integer.py index 70a029bd74bda..64a2d4f32aab3 100644 --- a/pandas/tests/arrays/test_integer.py +++ b/pandas/tests/arrays/test_integer.py @@ -793,6 +793,24 @@ def test_integer_array_constructor_copy(): assert result._mask is not mask +@pytest.mark.parametrize( + "data", + [ + pd.Series([1, 2], dtype="int64").tolist()[0], + pd.Series([1, 2], dtype="Int64").tolist()[0], + pd.Series([1, 2], dtype="int64").to_dict()[0], + pd.Series([1, 2], dtype="Int64").to_dict()[0], + list(pd.Series([1, 2], dtype="int64").iteritems())[0][1], + list(pd.Series([1, 2], dtype="Int64").iteritems())[0][1], + list(iter(pd.Series([1, 2], dtype="int64")))[0], + list(iter(pd.Series([1, 2], dtype="Int64")))[0], + ], +) +def test_integer_Series_iter_return_native(data): + # GH 29738 + assert isinstance(data, int) + + @pytest.mark.parametrize( "values", [