From 0f0a936a8ddaee26edcc2dce0f3913ec000885db Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Torsten=20W=C3=B6rtwein?= Date: Sat, 3 Sep 2022 20:01:40 -0400 Subject: [PATCH] TYP: contextmanager expects a Generator --- pandas/_config/config.py | 4 ++-- pandas/_config/localization.py | 4 ++-- pandas/_testing/_warnings.py | 3 ++- pandas/_testing/contexts.py | 18 ++++++++++-------- pandas/compat/pickle_compat.py | 4 ++-- pandas/core/common.py | 4 ++-- pandas/core/groupby/groupby.py | 3 ++- pandas/io/formats/format.py | 4 ++-- pandas/io/formats/style.py | 3 ++- pandas/plotting/_matplotlib/converter.py | 4 ++-- pandas/plotting/_misc.py | 4 ++-- pandas/tests/io/pytables/common.py | 5 ++++- pandas/tests/libs/test_hashtable.py | 3 ++- pandas/tests/test_register_accessor.py | 3 ++- pandas/util/_exceptions.py | 4 ++-- pandas/util/_test_decorators.py | 4 ++-- 16 files changed, 42 insertions(+), 32 deletions(-) diff --git a/pandas/_config/config.py b/pandas/_config/config.py index fc35b95bba7dd..e1a6cf04a435e 100644 --- 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 index c4355e954c67c..eaae30ee9098b 100644 --- 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 index a5b0d1e199863..f42004bdfdef3 100644 --- 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 index e64adb06bea7a..9647d46ac3391 100644 --- 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 index 813e8de72f96e..a9ae5f89b1a47 100644 --- 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 index 41ed68e73a4c0..e276e2d37d86c 100644 --- 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 index 89c9f3701a424..7cd6a93fea92d 100644 --- 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 index 27094fff5f812..2b388411e6de7 100644 --- 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 index a2a9079642344..d670931625a08 100644 --- 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 index 51a081645373e..5fe97534112b4 100644 --- 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 index 59117a3400bfd..4f197b5fc7095 100644 --- 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 index 67c3a2902dbcb..88a32e1a75972 100644 --- 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 index 0a3a10315b5fd..0cd340ab80897 100644 --- 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 index 3e4e57414269a..f0f330cc741b7 100644 --- 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 index fcd191e25ced5..ebd92e8366e18 100644 --- 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 index 5a799cdb89146..f94b2511da678 100644 --- 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. """