HDFStore.__iter__ (#539)

* HDFStore.__iter__

* enable Y034

* added missing IntervalIndex.closed

Author: twoertwein

.pre-commit-config.yaml

@@ -30,7 +30,7 @@ repos:
         - flake8-pyi==23.1.2
         types: [pyi]
         args: [
-          --ignore=E301 E302 E305 E402 E501 E701 E704 F401 F811 W503 Y019 Y034 Y037 Y041 Y042,
+          --ignore=E301 E302 E305 E402 E501 E701 E704 F401 F811 W503 Y019 Y037 Y041 Y042,
           # TypeVars in private files are already private
           --per-file-ignores=_*.pyi:Y001
         ]

pandas-stubs/_libs/missing.pyi

@@ -1,5 +1,7 @@
+from typing_extensions import Self
+
 class NAType:
-    def __new__(cls, *args, **kwargs) -> NAType: ...
+    def __new__(cls, *args, **kwargs) -> Self: ...
     def __format__(self, format_spec: str) -> str: ...
     def __bool__(self) -> None: ...
     def __hash__(self) -> int: ...

pandas-stubs/core/indexes/category.pyi

@@ -4,6 +4,7 @@ import numpy as np
 from pandas.core import accessor
 from pandas.core.indexes.base import Index  # , maybe_extract_name
 from pandas.core.indexes.extension import ExtensionIndex
+from typing_extensions import Self
 
 from pandas._typing import DtypeArg
 
@@ -18,7 +19,7 @@ class CategoricalIndex(ExtensionIndex, accessor.PandasDelegate):
         dtype=...,
         copy: bool = ...,
         name=...,
-    ) -> CategoricalIndex: ...
+    ) -> Self: ...
     def equals(self, other): ...
     @property
     def inferred_type(self) -> str: ...

pandas-stubs/core/indexes/interval.pyi

@@ -70,6 +70,8 @@ _TimestampLike: TypeAlias = Union[pd.Timestamp, np.datetime64, dt.datetime]
 _TimedeltaLike: TypeAlias = Union[pd.Timedelta, np.timedelta64, dt.timedelta]
 
 class IntervalIndex(IntervalMixin, Generic[IntervalT]):
+    closed: IntervalClosedType
+
     def __new__(
         cls,
         data: Sequence[IntervalT],

pandas-stubs/core/indexes/multi.pyi

@@ -11,6 +11,7 @@ from typing import (
 import numpy as np
 import pandas as pd
 from pandas.core.indexes.base import Index
+from typing_extensions import Self
 
 from pandas._typing import (
     T1,
@@ -32,7 +33,7 @@ class MultiIndex(Index):
         name=...,
         verify_integrity: bool = ...,
         _set_identity: bool = ...,
-    ) -> MultiIndex: ...
+    ) -> Self: ...
     def __init__(
         self,
         levels=...,

pandas-stubs/io/excel/_base.pyi

@@ -16,6 +16,7 @@ from odf.opendocument import OpenDocument
 from openpyxl.workbook.workbook import Workbook
 from pandas.core.frame import DataFrame
 import pyxlsb.workbook
+from typing_extensions import Self
 from xlrd.book import Book
 
 from pandas._typing import (
@@ -130,7 +131,7 @@ class ExcelWriter:
     @property
     def if_sheet_exists(self) -> Literal["error", "new", "replace", "overlay"]: ...
     def __fspath__(self) -> str: ...
-    def __enter__(self) -> ExcelWriter: ...
+    def __enter__(self) -> Self: ...
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
@@ -214,7 +215,7 @@ class ExcelFile:
     @property
     def sheet_names(self) -> list[int | str]: ...
     def close(self) -> None: ...
-    def __enter__(self) -> ExcelFile: ...
+    def __enter__(self) -> Self: ...
     def __exit__(
         self,
         exc_type: type[BaseException] | None,

pandas-stubs/io/parsers/readers.pyi

@@ -18,6 +18,7 @@ from typing import (
 from pandas.core.frame import DataFrame
 from pandas.core.indexes.base import Index
 from pandas.core.series import Series
+from typing_extensions import Self
 
 from pandas._typing import (
     CompressionOptions,
@@ -479,7 +480,7 @@ class TextFileReader(abc.Iterator):
     def read(self, nrows: int | None = ...) -> DataFrame: ...
     def get_chunk(self, size: int | None = ...) -> DataFrame: ...
     def __next__(self) -> DataFrame: ...
-    def __enter__(self) -> TextFileReader: ...
+    def __enter__(self) -> Self: ...
     def __exit__(
         self,
         exc_type: type[BaseException] | None,

pandas-stubs/io/pytables.pyi

@@ -1,5 +1,6 @@
 from collections.abc import (
     Generator,
+    Iterator,
     Sequence,
 )
 from types import TracebackType
@@ -15,6 +16,7 @@ from pandas import (
 )
 from pandas.core.computation.pytables import PyTablesExpr
 from pandas.core.generic import NDFrame
+from typing_extensions import Self
 
 from pandas._typing import (
     FilePath,
@@ -113,15 +115,15 @@ class HDFStore:
     def __getattr__(self, name: str) -> DataFrame | Series: ...
     def __contains__(self, key: str) -> bool: ...
     def __len__(self) -> int: ...
-    def __enter__(self) -> HDFStore: ...
+    def __enter__(self) -> Self: ...
     def __exit__(
         self,
         exc_type: type[BaseException] | None,
         exc_value: BaseException | None,
         traceback: TracebackType | None,
     ) -> None: ...
     def keys(self) -> list[str]: ...
-    def __iter__(self) -> list[str]: ...
+    def __iter__(self) -> Iterator[str]: ...
     def open(self, mode: Literal["a", "w", "r", "r+"] = ..., **kwargs) -> None: ...
     def close(self) -> None: ...
     @property

pandas-stubs/io/sas/sasreader.pyi

@@ -9,6 +9,7 @@ from typing import (
 )
 
 from pandas import DataFrame
+from typing_extensions import Self
 
 from pandas._typing import (
     CompressionOptions as CompressionOptions,
@@ -24,7 +25,7 @@ class ReaderBase(metaclass=ABCMeta):
     def read(self, nrows: int | None = ...) -> DataFrame: ...
     @abstractmethod
     def close(self) -> None: ...
-    def __enter__(self) -> ReaderBase: ...
+    def __enter__(self) -> Self: ...
     def __exit__(self, exc_type, exc_value, traceback) -> None: ...
 
 @overload

pandas-stubs/io/stata.pyi

@@ -9,6 +9,7 @@ from typing import (
 )
 
 from pandas.core.frame import DataFrame
+from typing_extensions import Self
 
 from pandas._typing import (
     CompressionOptions,
@@ -92,7 +93,7 @@ class StataReader(StataParser, abc.Iterator):
         compression: CompressionOptions = ...,
         storage_options: StorageOptions = ...,
     ) -> None: ...
-    def __enter__(self) -> StataReader: ...
+    def __enter__(self) -> Self: ...
     def __exit__(
         self,
         exc_type: type[BaseException] | None,

tests/test_interval_index.py

@@ -3,6 +3,8 @@
 import pandas as pd
 from typing_extensions import assert_type
 
+from pandas._typing import IntervalClosedType
+
 from tests import check
 
 
@@ -34,3 +36,5 @@ def test_from_tuples() -> None:
 def test_is_overlapping() -> None:
     ind = pd.IntervalIndex.from_tuples([(0, 2), (1, 3), (4, 5)])
     check(assert_type(ind.is_overlapping, bool), bool)
+
+    check(assert_type(ind.closed, IntervalClosedType), str)

tests/test_io.py

@@ -329,6 +329,8 @@ def test_hdfstore():
             DataFrame,
         )
         check(assert_type(store.get("df"), Union[DataFrame, Series]), DataFrame)
+        for key in store:
+            check(assert_type(key, str), str)
         check(assert_type(store.close(), None), type(None))
 
         store = HDFStore(path, model="r")