TYP: annotate plotting based on _get_axe_freq (#35960)

Author: jbrockmendel

pandas/plotting/_matplotlib/core.py

@@ -1,5 +1,5 @@
 import re
-from typing import List, Optional
+from typing import TYPE_CHECKING, List, Optional
 import warnings
 
 from matplotlib.artist import Artist
@@ -43,6 +43,9 @@
     table,
 )
 
+if TYPE_CHECKING:
+    from matplotlib.axes import Axes
+
 
 class MPLPlot:
     """
@@ -1147,7 +1150,7 @@ def _plot(cls, ax, x, y, style=None, column_num=None, stacking_id=None, **kwds):
         return lines
 
     @classmethod
-    def _ts_plot(cls, ax, x, data, style=None, **kwds):
+    def _ts_plot(cls, ax: "Axes", x, data, style=None, **kwds):
         from pandas.plotting._matplotlib.timeseries import (
             _decorate_axes,
             _maybe_resample,

pandas/plotting/_matplotlib/timeseries.py

@@ -24,14 +24,15 @@
 from pandas.tseries.frequencies import get_period_alias, is_subperiod, is_superperiod
 
 if TYPE_CHECKING:
-    from pandas import Index, Series  # noqa:F401
+    from matplotlib.axes import Axes
 
+    from pandas import Index, Series  # noqa:F401
 
 # ---------------------------------------------------------------------
 # Plotting functions and monkey patches
 
 
-def _maybe_resample(series: "Series", ax, kwargs):
+def _maybe_resample(series: "Series", ax: "Axes", kwargs):
     # resample against axes freq if necessary
     freq, ax_freq = _get_freq(ax, series)
 
@@ -74,7 +75,7 @@ def _is_sup(f1: str, f2: str) -> bool:
     )
 
 
-def _upsample_others(ax, freq, kwargs):
+def _upsample_others(ax: "Axes", freq, kwargs):
     legend = ax.get_legend()
     lines, labels = _replot_ax(ax, freq, kwargs)
     _replot_ax(ax, freq, kwargs)
@@ -97,7 +98,7 @@ def _upsample_others(ax, freq, kwargs):
         ax.legend(lines, labels, loc="best", title=title)
 
 
-def _replot_ax(ax, freq, kwargs):
+def _replot_ax(ax: "Axes", freq, kwargs):
     data = getattr(ax, "_plot_data", None)
 
     # clear current axes and data
@@ -127,7 +128,7 @@ def _replot_ax(ax, freq, kwargs):
     return lines, labels
 
 
-def _decorate_axes(ax, freq, kwargs):
+def _decorate_axes(ax: "Axes", freq, kwargs):
     """Initialize axes for time-series plotting"""
     if not hasattr(ax, "_plot_data"):
         ax._plot_data = []
@@ -143,7 +144,7 @@ def _decorate_axes(ax, freq, kwargs):
     ax.date_axis_info = None
 
 
-def _get_ax_freq(ax):
+def _get_ax_freq(ax: "Axes"):
     """
     Get the freq attribute of the ax object if set.
     Also checks shared axes (eg when using secondary yaxis, sharex=True
@@ -174,7 +175,7 @@ def _get_period_alias(freq) -> Optional[str]:
     return freq
 
 
-def _get_freq(ax, series: "Series"):
+def _get_freq(ax: "Axes", series: "Series"):
     # get frequency from data
     freq = getattr(series.index, "freq", None)
     if freq is None:
@@ -192,7 +193,7 @@ def _get_freq(ax, series: "Series"):
     return freq, ax_freq
 
 
-def _use_dynamic_x(ax, data: "FrameOrSeriesUnion") -> bool:
+def _use_dynamic_x(ax: "Axes", data: FrameOrSeriesUnion) -> bool:
     freq = _get_index_freq(data.index)
     ax_freq = _get_ax_freq(ax)
 
@@ -234,7 +235,7 @@ def _get_index_freq(index: "Index") -> Optional[BaseOffset]:
     return freq
 
 
-def _maybe_convert_index(ax, data):
+def _maybe_convert_index(ax: "Axes", data):
     # tsplot converts automatically, but don't want to convert index
     # over and over for DataFrames
     if isinstance(data.index, (ABCDatetimeIndex, ABCPeriodIndex)):
@@ -264,7 +265,7 @@ def _maybe_convert_index(ax, data):
 # Do we need the rest for convenience?
 
 
-def _format_coord(freq, t, y):
+def _format_coord(freq, t, y) -> str:
     time_period = Period(ordinal=int(t), freq=freq)
     return f"t = {time_period}  y = {y:8f}"