Annotate __iter__ method of Series (#126)

Author: diepala

pandas-stubs/core/generic.pyi

@@ -94,7 +94,6 @@ class NDFrame(PandasObject, indexing.IndexingMixin):
     def bool(self) -> _bool: ...
     def __abs__(self) -> NDFrame: ...
     def __round__(self, decimals: int = ...) -> NDFrame: ...
-    def __iter__(self) -> Iterator: ...
     def keys(self): ...
     def iteritems(self): ...
     def __len__(self) -> int: ...

pandas-stubs/core/series.pyi

@@ -10,6 +10,7 @@ from typing import (
     Generic,
     Hashable,
     Iterable,
+    Iterator,
     List,
     Literal,
     Mapping,
@@ -1772,6 +1773,7 @@ class Series(IndexOpsMixin, NDFrame, Generic[S1]):
     def set_axis(
         self, labels, axis: Axis = ..., inplace: bool = ...
     ) -> Optional[Series[S1]]: ...
+    def __iter__(self) -> Iterator[S1]: ...
 
 class TimestampSeries(Series): ...
 

tests/test_series.py

@@ -2,6 +2,8 @@
 import tempfile
 from typing import (
     TYPE_CHECKING,
+    Iterable,
+    Iterator,
     List,
     Sequence,
     Union,
@@ -787,3 +789,10 @@ def test_iloc_setitem_ndarray() -> None:
     values_s.iloc[indices_u16] = -1
     values_s.iloc[indices_u32] = -1
     values_s.iloc[indices_u64] = -1
+
+
+def test_types_iter() -> None:
+    s = pd.Series([1, 2, 3], dtype=int)
+    iterable: Iterable[int] = s
+    assert_type(iter(s), Iterator[int])
+    assert_type(next(iter(s)), int)