Fix polar coordinate type checking

Author: iangow

plotnine/coords/coord_polar.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, cast
 
 import numpy as np
 
@@ -11,6 +11,7 @@
 if TYPE_CHECKING:
     import pandas as pd
     from matplotlib.axes import Axes
+    from matplotlib.projections.polar import PolarAxes
 
     from plotnine.iapi import panel_view
     from plotnine.scales.scale import scale
@@ -216,8 +217,8 @@ def backtransform_range(self, panel_params: panel_view) -> panel_ranges:
         t_range = tuple(self.params["theta_range"])
         r_range = tuple(self.params["r_range"])
         if self.theta == "x":
-            return panel_ranges(x=t_range, y=r_range)  # type: ignore[arg-type]
-        return panel_ranges(x=r_range, y=t_range)  # type: ignore[arg-type]
+            return panel_ranges(x=t_range, y=r_range)
+        return panel_ranges(x=r_range, y=t_range)
 
     # ------------------------------------------------------------------
     # Draw decorations on PolarAxes
@@ -231,10 +232,11 @@ def draw(self, axs: list[Axes]) -> None:
         mpl_direction = -1 if self.direction == 1 else 1
 
         for ax in axs:
-            ax.set_theta_zero_location("N")  # 12 o'clock = 0
-            ax.set_theta_direction(mpl_direction)
+            polar_ax = cast("PolarAxes", ax)
+            polar_ax.set_theta_zero_location("N")  # 12 o'clock = 0
+            polar_ax.set_theta_direction(mpl_direction)
             if np.isfinite(r_min) and np.isfinite(r_max) and r_min != r_max:
-                ax.set_rlim(float(r_min), float(r_max))
+                polar_ax.set_rlim(float(r_min), float(r_max))
 
     # ------------------------------------------------------------------
     # Misc

plotnine/coords/coord_radial.py

@@ -1,7 +1,7 @@
 from __future__ import annotations
 
 from dataclasses import replace
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Sequence, cast
 
 import numpy as np
 
@@ -10,6 +10,7 @@
 if TYPE_CHECKING:
     import pandas as pd
     from matplotlib.axes import Axes
+    from matplotlib.projections.polar import PolarAxes
 
     from plotnine.iapi import panel_view
     from plotnine.scales.scale import scale
@@ -184,7 +185,8 @@ def setup_panel_params(self, scale_x: scale, scale_y: scale) -> panel_view:
         if self.rlim is not None:
             self.params["r_range"] = tuple(self.rlim)
             rlo, rhi = self.rlim
-            breaks, labels = pv.y.breaks, pv.y.labels
+            breaks = cast("Sequence[float]", pv.y.breaks)
+            labels = pv.y.labels
             mask = [rlo <= b <= rhi for b in breaks]
             new_y = replace(
                 pv.y,
@@ -287,11 +289,12 @@ def draw(self, axs: list[Axes]) -> None:
         arc = self._arc
 
         for ax in axs:
+            polar_ax = cast("PolarAxes", ax)
             # Restrict visible theta range for partial arcs.
             if self.end is not None:
                 theta_lo = min(self.start, self.start + arc)
                 theta_hi = max(self.start, self.start + arc)
-                ax.set_thetalim(theta_lo, theta_hi)
+                polar_ax.set_thetalim(theta_lo, theta_hi)
 
             # Inner radius: push the data away from the centre by setting a
             # virtual r-origin below r_min.  Formula: solve
@@ -306,16 +309,18 @@ def draw(self, axs: list[Axes]) -> None:
                 r_origin = (r_min - self.inner_radius * r_max) / (
                     1.0 - self.inner_radius
                 )
-                ax.set_rorigin(r_origin)
+                polar_ax.set_rorigin(r_origin)
 
             # Radial axis label placement.
             if self.r_axis_inside is not None:
                 if isinstance(self.r_axis_inside, bool):
                     if self.r_axis_inside:
                         # Just inside the start angle keeps it out of the data.
-                        ax.set_rlabel_position(np.degrees(self.start) + 10)
+                        polar_ax.set_rlabel_position(
+                            np.degrees(self.start) + 10
+                        )
                 else:
-                    ax.set_rlabel_position(
+                    polar_ax.set_rlabel_position(
                         np.degrees(float(self.r_axis_inside))
                     )