TYP: contextmanager expects a Generator (#48383)

twoertwein · web-flow · commit 6f8ab49ce312 · 2022-09-06T13:14:16.000-07:00
diff --git a/pandas/_config/config.py b/pandas/_config/config.py
@@ -59,9 +59,9 @@
 from typing import (
     Any,
     Callable,
+    Generator,
     Generic,
     Iterable,
-    Iterator,
     NamedTuple,
     cast,
 )
@@ -743,7 +743,7 @@ def pp(name: str, ks: Iterable[str]) -> list[str]:
 
 
 @contextmanager
-def config_prefix(prefix) -> Iterator[None]:
+def config_prefix(prefix) -> Generator[None, None, None]:
     """
     contextmanager for multiple invocations of API with a common prefix
 
diff --git a/pandas/_config/localization.py b/pandas/_config/localization.py
@@ -11,7 +11,7 @@
 import subprocess
 from typing import (
     Callable,
-    Iterator,
+    Generator,
 )
 
 from pandas._config.config import options
@@ -20,7 +20,7 @@
 @contextmanager
 def set_locale(
     new_locale: str | tuple[str, str], lc_var: int = locale.LC_ALL
-) -> Iterator[str | tuple[str, str]]:
+) -> Generator[str | tuple[str, str], None, None]:
     """
     Context manager for temporarily setting a locale.
 
diff --git a/pandas/_testing/_warnings.py b/pandas/_testing/_warnings.py
@@ -7,6 +7,7 @@
 import re
 import sys
 from typing import (
+    Generator,
     Literal,
     Sequence,
     Type,
@@ -24,7 +25,7 @@ def assert_produces_warning(
     check_stacklevel: bool = True,
     raise_on_extra_warnings: bool = True,
     match: str | None = None,
-):
+) -> Generator[list[warnings.WarningMessage], None, None]:
     """
     Context manager for running code expected to either raise a specific warning,
     multiple specific warnings, or not raise any warnings. Verifies that the code
diff --git a/pandas/_testing/contexts.py b/pandas/_testing/contexts.py
@@ -8,7 +8,7 @@
 from typing import (
     IO,
     Any,
-    Iterator,
+    Generator,
 )
 import uuid
 
@@ -20,7 +20,7 @@
 
 
 @contextmanager
-def decompress_file(path, compression) -> Iterator[IO[bytes]]:
+def decompress_file(path, compression) -> Generator[IO[bytes], None, None]:
     """
     Open a compressed file and return a file object.
 
@@ -41,7 +41,7 @@ def decompress_file(path, compression) -> Iterator[IO[bytes]]:
 
 
 @contextmanager
-def set_timezone(tz: str) -> Iterator[None]:
+def set_timezone(tz: str) -> Generator[None, None, None]:
     """
     Context manager for temporarily setting a timezone.
 
@@ -84,7 +84,9 @@ def setTZ(tz):
 
 
 @contextmanager
-def ensure_clean(filename=None, return_filelike: bool = False, **kwargs: Any):
+def ensure_clean(
+    filename=None, return_filelike: bool = False, **kwargs: Any
+) -> Generator[Any, None, None]:
     """
     Gets a temporary path and agrees to remove on close.
 
@@ -127,7 +129,7 @@ def ensure_clean(filename=None, return_filelike: bool = False, **kwargs: Any):
 
 
 @contextmanager
-def ensure_clean_dir() -> Iterator[str]:
+def ensure_clean_dir() -> Generator[str, None, None]:
     """
     Get a temporary directory path and agrees to remove on close.
 
@@ -146,7 +148,7 @@ def ensure_clean_dir() -> Iterator[str]:
 
 
 @contextmanager
-def ensure_safe_environment_variables() -> Iterator[None]:
+def ensure_safe_environment_variables() -> Generator[None, None, None]:
     """
     Get a context manager to safely set environment variables
 
@@ -162,7 +164,7 @@ def ensure_safe_environment_variables() -> Iterator[None]:
 
 
 @contextmanager
-def with_csv_dialect(name, **kwargs) -> Iterator[None]:
+def with_csv_dialect(name, **kwargs) -> Generator[None, None, None]:
     """
     Context manager to temporarily register a CSV dialect for parsing CSV.
 
@@ -196,7 +198,7 @@ def with_csv_dialect(name, **kwargs) -> Iterator[None]:
 
 
 @contextmanager
-def use_numexpr(use, min_elements=None) -> Iterator[None]:
+def use_numexpr(use, min_elements=None) -> Generator[None, None, None]:
     from pandas.core.computation import expressions as expr
 
     if min_elements is None:
diff --git a/pandas/compat/pickle_compat.py b/pandas/compat/pickle_compat.py
@@ -9,7 +9,7 @@
 import pickle as pkl
 from typing import (
     TYPE_CHECKING,
-    Iterator,
+    Generator,
 )
 import warnings
 
@@ -294,7 +294,7 @@ def loads(
 
 
 @contextlib.contextmanager
-def patch_pickle() -> Iterator[None]:
+def patch_pickle() -> Generator[None, None, None]:
     """
     Temporarily patch pickle to use our unpickler.
     """
diff --git a/pandas/core/common.py b/pandas/core/common.py
@@ -18,9 +18,9 @@
     Any,
     Callable,
     Collection,
+    Generator,
     Hashable,
     Iterable,
-    Iterator,
     Sequence,
     cast,
     overload,
@@ -534,7 +534,7 @@ def convert_to_list_like(
 
 
 @contextlib.contextmanager
-def temp_setattr(obj, attr: str, value) -> Iterator[None]:
+def temp_setattr(obj, attr: str, value) -> Generator[None, None, None]:
     """Temporarily set attribute on an object.
 
     Args:
diff --git a/pandas/core/groupby/groupby.py b/pandas/core/groupby/groupby.py
@@ -20,6 +20,7 @@ class providing the base-class of operations.
 from typing import (
     TYPE_CHECKING,
     Callable,
+    Generator,
     Hashable,
     Iterable,
     Iterator,
@@ -1106,7 +1107,7 @@ def _reset_group_selection(self) -> None:
             self._reset_cache("_selected_obj")
 
     @contextmanager
-    def _group_selection_context(self) -> Iterator[GroupBy]:
+    def _group_selection_context(self) -> Generator[GroupBy, None, None]:
         """
         Set / reset the _group_selection_context.
         """
diff --git a/pandas/io/formats/format.py b/pandas/io/formats/format.py
@@ -21,9 +21,9 @@
     Any,
     Callable,
     Final,
+    Generator,
     Hashable,
     Iterable,
-    Iterator,
     List,
     Mapping,
     Sequence,
@@ -1216,7 +1216,7 @@ def save_to_buffer(
 @contextmanager
 def get_buffer(
     buf: FilePath | WriteBuffer[str] | None, encoding: str | None = None
-) -> Iterator[WriteBuffer[str]] | Iterator[StringIO]:
+) -> Generator[WriteBuffer[str], None, None] | Generator[StringIO, None, None]:
     """
     Context manager to open, yield and close buffer for filenames or Path-like
     objects, otherwise yield buf unchanged.
diff --git a/pandas/io/formats/style.py b/pandas/io/formats/style.py
@@ -11,6 +11,7 @@
 from typing import (
     Any,
     Callable,
+    Generator,
     Hashable,
     Sequence,
     overload,
@@ -80,7 +81,7 @@
 
 
 @contextmanager
-def _mpl(func: Callable):
+def _mpl(func: Callable) -> Generator[tuple[Any, Any], None, None]:
     if has_mpl:
         yield plt, mpl
     else:
diff --git a/pandas/plotting/_matplotlib/converter.py b/pandas/plotting/_matplotlib/converter.py
@@ -12,7 +12,7 @@
     TYPE_CHECKING,
     Any,
     Final,
-    Iterator,
+    Generator,
     cast,
 )
 
@@ -99,7 +99,7 @@ def wrapper(*args, **kwargs):
 
 
 @contextlib.contextmanager
-def pandas_converters() -> Iterator[None]:
+def pandas_converters() -> Generator[None, None, None]:
     """
     Context manager registering pandas' converters for a plot.
 
diff --git a/pandas/plotting/_misc.py b/pandas/plotting/_misc.py
@@ -3,7 +3,7 @@
 from contextlib import contextmanager
 from typing import (
     TYPE_CHECKING,
-    Iterator,
+    Generator,
 )
 
 from pandas.plotting._core import _get_plot_backend
@@ -594,7 +594,7 @@ def _get_canonical_key(self, key):
         return self._ALIASES.get(key, key)
 
     @contextmanager
-    def use(self, key, value) -> Iterator[_Options]:
+    def use(self, key, value) -> Generator[_Options, None, None]:
         """
         Temporarily set a parameter value using the with statement.
         Aliasing allowed.
diff --git a/pandas/tests/io/pytables/common.py b/pandas/tests/io/pytables/common.py
@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 import os
 import tempfile
+from typing import Generator
 
 import pytest
 
@@ -36,7 +37,9 @@ def create_tempfile(path):
 
 # contextmanager to ensure the file cleanup
 @contextmanager
-def ensure_clean_store(path, mode="a", complevel=None, complib=None, fletcher32=False):
+def ensure_clean_store(
+    path, mode="a", complevel=None, complib=None, fletcher32=False
+) -> Generator[HDFStore, None, None]:
 
     try:
 
diff --git a/pandas/tests/libs/test_hashtable.py b/pandas/tests/libs/test_hashtable.py
@@ -1,6 +1,7 @@
 from contextlib import contextmanager
 import struct
 import tracemalloc
+from typing import Generator
 
 import numpy as np
 import pytest
@@ -13,7 +14,7 @@
 
 
 @contextmanager
-def activated_tracemalloc():
+def activated_tracemalloc() -> Generator[None, None, None]:
     tracemalloc.start()
     try:
         yield
diff --git a/pandas/tests/test_register_accessor.py b/pandas/tests/test_register_accessor.py
@@ -1,4 +1,5 @@
 import contextlib
+from typing import Generator
 
 import pytest
 
@@ -23,7 +24,7 @@ def __init__(self) -> None:
 
 
 @contextlib.contextmanager
-def ensure_removed(obj, attr):
+def ensure_removed(obj, attr) -> Generator[None, None, None]:
     """Ensure that an attribute added to 'obj' during the test is
     removed when we're done
     """
diff --git a/pandas/util/_exceptions.py b/pandas/util/_exceptions.py
@@ -4,11 +4,11 @@
 import functools
 import inspect
 import os
-from typing import Iterator
+from typing import Generator
 
 
 @contextlib.contextmanager
-def rewrite_exception(old_name: str, new_name: str) -> Iterator[None]:
+def rewrite_exception(old_name: str, new_name: str) -> Generator[None, None, None]:
     """
     Rewrite the message of an exception.
     """
diff --git a/pandas/util/_test_decorators.py b/pandas/util/_test_decorators.py
@@ -30,7 +30,7 @@ def test_foo():
 import locale
 from typing import (
     Callable,
-    Iterator,
+    Generator,
 )
 import warnings
 
@@ -257,7 +257,7 @@ def check_file_leaks(func) -> Callable:
 
 
 @contextmanager
-def file_leak_context() -> Iterator[None]:
+def file_leak_context() -> Generator[None, None, None]:
     """
     ContextManager analogue to check_file_leaks.
     """
