TYP: Missing return annotations in util/tseries/plotting (#47510)

* TYP: Missing return annotations in util/tseries/plotting

* the more tricky parts

Author: twoertwein

pandas/plotting/_misc.py

@@ -1,7 +1,18 @@
+from __future__ import annotations
+
 from contextlib import contextmanager
+from typing import (
+    TYPE_CHECKING,
+    Iterator,
+)
 
 from pandas.plotting._core import _get_plot_backend
 
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+    from matplotlib.figure import Figure
+    import numpy as np
+
 
 def table(ax, data, rowLabels=None, colLabels=None, **kwargs):
     """
@@ -27,7 +38,7 @@ def table(ax, data, rowLabels=None, colLabels=None, **kwargs):
     )
 
 
-def register():
+def register() -> None:
     """
     Register pandas formatters and converters with matplotlib.
 
@@ -49,7 +60,7 @@ def register():
     plot_backend.register()
 
 
-def deregister():
+def deregister() -> None:
     """
     Remove pandas formatters and converters.
 
@@ -81,7 +92,7 @@ def scatter_matrix(
     hist_kwds=None,
     range_padding=0.05,
     **kwargs,
-):
+) -> np.ndarray:
     """
     Draw a matrix of scatter plots.
 
@@ -156,7 +167,7 @@ def scatter_matrix(
     )
 
 
-def radviz(frame, class_column, ax=None, color=None, colormap=None, **kwds):
+def radviz(frame, class_column, ax=None, color=None, colormap=None, **kwds) -> Axes:
     """
     Plot a multidimensional dataset in 2D.
 
@@ -239,7 +250,7 @@ def radviz(frame, class_column, ax=None, color=None, colormap=None, **kwds):
 
 def andrews_curves(
     frame, class_column, ax=None, samples=200, color=None, colormap=None, **kwargs
-):
+) -> Axes:
     """
     Generate a matplotlib plot of Andrews curves, for visualising clusters of
     multivariate data.
@@ -297,7 +308,7 @@ def andrews_curves(
     )
 
 
-def bootstrap_plot(series, fig=None, size=50, samples=500, **kwds):
+def bootstrap_plot(series, fig=None, size=50, samples=500, **kwds) -> Figure:
     """
     Bootstrap plot on mean, median and mid-range statistics.
 
@@ -364,7 +375,7 @@ def parallel_coordinates(
     axvlines_kwds=None,
     sort_labels=False,
     **kwargs,
-):
+) -> Axes:
     """
     Parallel coordinates plotting.
 
@@ -430,7 +441,7 @@ def parallel_coordinates(
     )
 
 
-def lag_plot(series, lag=1, ax=None, **kwds):
+def lag_plot(series, lag=1, ax=None, **kwds) -> Axes:
     """
     Lag plot for time series.
 
@@ -474,7 +485,7 @@ def lag_plot(series, lag=1, ax=None, **kwds):
     return plot_backend.lag_plot(series=series, lag=lag, ax=ax, **kwds)
 
 
-def autocorrelation_plot(series, ax=None, **kwargs):
+def autocorrelation_plot(series, ax=None, **kwargs) -> Axes:
     """
     Autocorrelation plot for time series.
 
@@ -531,21 +542,21 @@ def __getitem__(self, key):
             raise ValueError(f"{key} is not a valid pandas plotting option")
         return super().__getitem__(key)
 
-    def __setitem__(self, key, value):
+    def __setitem__(self, key, value) -> None:
         key = self._get_canonical_key(key)
-        return super().__setitem__(key, value)
+        super().__setitem__(key, value)
 
-    def __delitem__(self, key):
+    def __delitem__(self, key) -> None:
         key = self._get_canonical_key(key)
         if key in self._DEFAULT_KEYS:
             raise ValueError(f"Cannot remove default parameter {key}")
-        return super().__delitem__(key)
+        super().__delitem__(key)
 
     def __contains__(self, key) -> bool:
         key = self._get_canonical_key(key)
         return super().__contains__(key)
 
-    def reset(self):
+    def reset(self) -> None:
         """
         Reset the option store to its initial state
 
@@ -560,7 +571,7 @@ def _get_canonical_key(self, key):
         return self._ALIASES.get(key, key)
 
     @contextmanager
-    def use(self, key, value):
+    def use(self, key, value) -> Iterator[_Options]:
         """
         Temporarily set a parameter value using the with statement.
         Aliasing allowed.

pandas/tseries/frequencies.py

@@ -314,12 +314,12 @@ def get_freq(self) -> str | None:
             return _maybe_add_count("N", delta)
 
     @cache_readonly
-    def day_deltas(self):
+    def day_deltas(self) -> list[int]:
         ppd = periods_per_day(self._reso)
         return [x / ppd for x in self.deltas]
 
     @cache_readonly
-    def hour_deltas(self):
+    def hour_deltas(self) -> list[int]:
         pph = periods_per_day(self._reso) // 24
         return [x / pph for x in self.deltas]
 
@@ -328,10 +328,10 @@ def fields(self) -> np.ndarray:  # structured array of fields
         return build_field_sarray(self.i8values, reso=self._reso)
 
     @cache_readonly
-    def rep_stamp(self):
+    def rep_stamp(self) -> Timestamp:
         return Timestamp(self.i8values[0])
 
-    def month_position_check(self):
+    def month_position_check(self) -> str | None:
         return month_position_check(self.fields, self.index.dayofweek)
 
     @cache_readonly
@@ -394,7 +394,11 @@ def _get_annual_rule(self) -> str | None:
             return None
 
         pos_check = self.month_position_check()
-        return {"cs": "AS", "bs": "BAS", "ce": "A", "be": "BA"}.get(pos_check)
+        #  error: Argument 1 to "get" of "dict" has incompatible type
+        # "Optional[str]"; expected "str"
+        return {"cs": "AS", "bs": "BAS", "ce": "A", "be": "BA"}.get(
+            pos_check  # type: ignore[arg-type]
+        )
 
     def _get_quarterly_rule(self) -> str | None:
         if len(self.mdiffs) > 1:
@@ -404,13 +408,21 @@ def _get_quarterly_rule(self) -> str | None:
             return None
 
         pos_check = self.month_position_check()
-        return {"cs": "QS", "bs": "BQS", "ce": "Q", "be": "BQ"}.get(pos_check)
+        # error: Argument 1 to "get" of "dict" has incompatible type
+        # "Optional[str]"; expected "str"
+        return {"cs": "QS", "bs": "BQS", "ce": "Q", "be": "BQ"}.get(
+            pos_check  # type: ignore[arg-type]
+        )
 
     def _get_monthly_rule(self) -> str | None:
         if len(self.mdiffs) > 1:
             return None
         pos_check = self.month_position_check()
-        return {"cs": "MS", "bs": "BMS", "ce": "M", "be": "BM"}.get(pos_check)
+        # error: Argument 1 to "get" of "dict" has incompatible type
+        # "Optional[str]"; expected "str"
+        return {"cs": "MS", "bs": "BMS", "ce": "M", "be": "BM"}.get(
+            pos_check  # type: ignore[arg-type]
+        )
 
     def _is_business_daily(self) -> bool:
         # quick check: cannot be business daily

pandas/util/_exceptions.py

@@ -3,10 +3,11 @@
 import contextlib
 import inspect
 import os
+from typing import Iterator
 
 
 @contextlib.contextmanager
-def rewrite_exception(old_name: str, new_name: str):
+def rewrite_exception(old_name: str, new_name: str) -> Iterator[None]:
     """
     Rewrite the message of an exception.
     """

pandas/util/_test_decorators.py

@@ -27,7 +27,10 @@ def test_foo():
 
 from contextlib import contextmanager
 import locale
-from typing import Callable
+from typing import (
+    Callable,
+    Iterator,
+)
 import warnings
 
 import numpy as np
@@ -253,7 +256,7 @@ def check_file_leaks(func) -> Callable:
 
 
 @contextmanager
-def file_leak_context():
+def file_leak_context() -> Iterator[None]:
     """
     ContextManager analogue to check_file_leaks.
     """
@@ -290,7 +293,7 @@ def async_mark():
     return async_mark
 
 
-def mark_array_manager_not_yet_implemented(request):
+def mark_array_manager_not_yet_implemented(request) -> None:
     mark = pytest.mark.xfail(reason="Not yet implemented for ArrayManager")
     request.node.add_marker(mark)
 

pandas/util/_validators.py

@@ -7,6 +7,8 @@
 from typing import (
     Iterable,
     Sequence,
+    TypeVar,
+    overload,
 )
 import warnings
 
@@ -19,6 +21,9 @@
     is_integer,
 )
 
+BoolishT = TypeVar("BoolishT", bool, int)
+BoolishNoneT = TypeVar("BoolishNoneT", bool, int, None)
+
 
 def _check_arg_length(fname, args, max_fname_arg_count, compat_args):
     """
@@ -78,7 +83,7 @@ def _check_for_default_values(fname, arg_val_dict, compat_args):
             )
 
 
-def validate_args(fname, args, max_fname_arg_count, compat_args):
+def validate_args(fname, args, max_fname_arg_count, compat_args) -> None:
     """
     Checks whether the length of the `*args` argument passed into a function
     has at most `len(compat_args)` arguments and whether or not all of these
@@ -132,7 +137,7 @@ def _check_for_invalid_keys(fname, kwargs, compat_args):
         raise TypeError(f"{fname}() got an unexpected keyword argument '{bad_arg}'")
 
 
-def validate_kwargs(fname, kwargs, compat_args):
+def validate_kwargs(fname, kwargs, compat_args) -> None:
     """
     Checks whether parameters passed to the **kwargs argument in a
     function `fname` are valid parameters as specified in `*compat_args`
@@ -159,7 +164,9 @@ def validate_kwargs(fname, kwargs, compat_args):
     _check_for_default_values(fname, kwds, compat_args)
 
 
-def validate_args_and_kwargs(fname, args, kwargs, max_fname_arg_count, compat_args):
+def validate_args_and_kwargs(
+    fname, args, kwargs, max_fname_arg_count, compat_args
+) -> None:
     """
     Checks whether parameters passed to the *args and **kwargs argument in a
     function `fname` are valid parameters as specified in `*compat_args`
@@ -215,7 +222,9 @@ def validate_args_and_kwargs(fname, args, kwargs, max_fname_arg_count, compat_ar
     validate_kwargs(fname, kwargs, compat_args)
 
 
-def validate_bool_kwarg(value, arg_name, none_allowed=True, int_allowed=False):
+def validate_bool_kwarg(
+    value: BoolishNoneT, arg_name, none_allowed=True, int_allowed=False
+) -> BoolishNoneT:
     """
     Ensure that argument passed in arg_name can be interpreted as boolean.
 
@@ -424,12 +433,22 @@ def validate_percentile(q: float | Iterable[float]) -> np.ndarray:
     return q_arr
 
 
+@overload
+def validate_ascending(ascending: BoolishT) -> BoolishT:
+    ...
+
+
+@overload
+def validate_ascending(ascending: Sequence[BoolishT]) -> list[BoolishT]:
+    ...
+
+
 def validate_ascending(
-    ascending: bool | int | Sequence[bool | int] = True,
-):
+    ascending: bool | int | Sequence[BoolishT],
+) -> bool | int | list[BoolishT]:
     """Validate ``ascending`` kwargs for ``sort_index`` method."""
     kwargs = {"none_allowed": False, "int_allowed": True}
-    if not isinstance(ascending, (list, tuple)):
+    if not isinstance(ascending, Sequence):
         return validate_bool_kwarg(ascending, "ascending", **kwargs)
 
     return [validate_bool_kwarg(item, "ascending", **kwargs) for item in ascending]