Move axes setup to coord classes

Author: iangow

plotnine/coords/coord.py

@@ -10,10 +10,11 @@
 if typing.TYPE_CHECKING:
     from typing import Any
 
+    from matplotlib.axes import Axes
     import numpy.typing as npt
     import pandas as pd
 
-    from plotnine import ggplot
+    from plotnine import ggplot, theme
     from plotnine.iapi import labels_view, panel_view
     from plotnine.scales.scale import scale
     from plotnine.typing import (
@@ -115,13 +116,50 @@ def draw(self, axs: list) -> None:
         add elements such as polar grid lines.
         """
 
-    def post_setup_ax(self, ax: Any) -> None:
+    def setup_ax(
+        self, ax: Axes, panel_params: panel_view, theme: theme
+    ) -> None:
         """
-        Hook called for each axes after set_limits_breaks_and_labels.
+        Set limits, breaks and labels for one panel axes.
 
-        Override in subclasses to apply per-axes settings that must
-        run after the facet has set tick positions and label padding.
+        Subclasses can override this to customize axes setup, or call
+        `super().setup_ax(...)` and add coordinate-specific behavior.
         """
+        from .._mpl.ticker import MyFixedFormatter
+
+        def _inf_to_none(
+            t: tuple[float, float],
+        ) -> tuple[float | None, float | None]:
+            """
+            Replace infinities with None
+            """
+            a = t[0] if np.isfinite(t[0]) else None
+            b = t[1] if np.isfinite(t[1]) else None
+            return (a, b)
+
+        # limits
+        ax.set_xlim(*_inf_to_none(panel_params.x.range))
+        ax.set_ylim(*_inf_to_none(panel_params.y.range))
+
+        # breaks, labels
+        ax.set_xticks(panel_params.x.breaks, panel_params.x.labels)
+        ax.set_yticks(panel_params.y.breaks, panel_params.y.labels)
+
+        # minor breaks
+        ax.set_xticks(panel_params.x.minor_breaks, minor=True)
+        ax.set_yticks(panel_params.y.minor_breaks, minor=True)
+
+        # When you manually set the tick labels MPL changes the locator
+        # so that it no longer reports the x & y positions
+        # Fixes https://github.com/has2k1/plotnine/issues/187
+        ax.xaxis.set_major_formatter(MyFixedFormatter(panel_params.x.labels))
+        ax.yaxis.set_major_formatter(MyFixedFormatter(panel_params.y.labels))
+
+        pad_x = theme.get_margin("axis_text_x").pt.t
+        pad_y = theme.get_margin("axis_text_y").pt.r
+
+        ax.tick_params(axis="x", which="major", pad=pad_x)
+        ax.tick_params(axis="y", which="major", pad=pad_y)
 
     def labels(self, cur_labels: labels_view) -> labels_view:
         """

plotnine/coords/coord_radial.py

@@ -12,6 +12,7 @@
     from matplotlib.axes import Axes
     from matplotlib.projections.polar import PolarAxes
 
+    from plotnine import theme
     from plotnine.iapi import panel_view
     from plotnine.scales.scale import scale
 
@@ -221,8 +222,8 @@ def setup_panel_params(self, scale_x: scale, scale_y: scale) -> panel_view:
             x_updates["labels"] = theta_labels
 
         # Partial arc: x panel range must match [arc_lo, arc_hi] so that
-        # set_limits_breaks_and_labels calls ax.set_xlim(arc_lo, arc_hi) rather
-        # than ax.set_xlim(0, 2π), which would override set_thetalim.
+        # coord.setup_ax calls ax.set_xlim(arc_lo, arc_hi) rather than
+        # ax.set_xlim(0, 2π), which would override set_thetalim.
         if arc_lo is not None:
             x_updates["limits"] = (arc_lo, arc_hi)
             x_updates["range"] = (arc_lo, arc_hi)
@@ -324,10 +325,13 @@ def draw(self, axs: list[Axes]) -> None:
                         np.degrees(float(self.r_axis_inside))
                     )
 
-    def post_setup_ax(self, ax: Axes) -> None:
+    def setup_ax(
+        self, ax: Axes, panel_params: panel_view, theme: theme
+    ) -> None:
         """
-        Apply theta label pad after facet has set tick positions and padding.
+        Apply theta label pad after setting tick positions and padding.
         """
+        super().setup_ax(ax, panel_params, theme)
         if self.theta_labels or self.end is not None:
             ax.tick_params(axis="x", pad=self.theta_label_pad)
         # Allow geom_text labels to extend past the polar axes bounding box

plotnine/facets/facet.py

@@ -26,7 +26,7 @@
     from plotnine.coords.coord import coord
     from plotnine.facets.labelling import CanBeStripLabellingFunc
     from plotnine.facets.layout import Layout
-    from plotnine.iapi import layout_details, panel_view
+    from plotnine.iapi import layout_details
     from plotnine.layer import Layers
     from plotnine.mapping import Environment
     from plotnine.scales.scale import scale
@@ -303,60 +303,6 @@ def make_strips(self, layout_info: layout_details, ax: Axes) -> Strips:
         """
         return Strips()
 
-    def set_limits_breaks_and_labels(self, panel_params: panel_view, ax: Axes):
-        """
-        Add limits, breaks and labels to the axes
-
-        Parameters
-        ----------
-        panel_params :
-            range information for the axes
-        ax :
-            Axes
-        """
-        from .._mpl.ticker import MyFixedFormatter
-
-        def _inf_to_none(
-            t: tuple[float, float],
-        ) -> tuple[float | None, float | None]:
-            """
-            Replace infinities with None
-            """
-            a = t[0] if np.isfinite(t[0]) else None
-            b = t[1] if np.isfinite(t[1]) else None
-            return (a, b)
-
-        theme = self.theme
-
-        # limits
-        ax.set_xlim(*_inf_to_none(panel_params.x.range))
-        ax.set_ylim(*_inf_to_none(panel_params.y.range))
-
-        if typing.TYPE_CHECKING:
-            assert callable(ax.set_xticks)
-            assert callable(ax.set_yticks)
-
-        # breaks, labels
-        ax.set_xticks(panel_params.x.breaks, panel_params.x.labels)
-        ax.set_yticks(panel_params.y.breaks, panel_params.y.labels)
-
-        # minor breaks
-        ax.set_xticks(panel_params.x.minor_breaks, minor=True)
-        ax.set_yticks(panel_params.y.minor_breaks, minor=True)
-
-        # When you manually set the tick labels MPL changes the locator
-        # so that it no longer reports the x & y positions
-        # Fixes https://github.com/has2k1/plotnine/issues/187
-        ax.xaxis.set_major_formatter(MyFixedFormatter(panel_params.x.labels))
-        ax.yaxis.set_major_formatter(MyFixedFormatter(panel_params.y.labels))
-
-        pad_x = theme.get_margin("axis_text_x").pt.t
-        pad_y = theme.get_margin("axis_text_y").pt.r
-
-        ax.tick_params(axis="x", which="major", pad=pad_x)
-        ax.tick_params(axis="y", which="major", pad=pad_y)
-        self.coordinates.post_setup_ax(ax)
-
     def __deepcopy__(self, memo: dict[Any, Any]) -> facet:
         """
         Deep copy without copying the dataframe and environment

plotnine/ggplot.py

@@ -567,7 +567,7 @@ def _draw_breaks_and_labels(self):
             pidx = layout_info.panel_index
             ax = self.axs[pidx]
             panel_params = self.layout.panel_params[pidx]
-            self.facet.set_limits_breaks_and_labels(panel_params, ax)
+            self.coordinates.setup_ax(ax, panel_params, self.theme)
 
             # Remove unnecessary ticks and labels
             if not layout_info.axis_x:

tests/test_coord_polar.py

@@ -10,6 +10,7 @@
     element_line,
     geom_col,
     geom_point,
+    geom_text,
     ggplot,
     theme,
 )
@@ -291,14 +292,18 @@ def test_coord_radial_draw_float_r_axis_position():
         plt.close(fig)
 
 
-def test_coord_radial_post_setup_ax_sets_pad_and_unclips_text():
-    coord = coord_radial(theta_label_pad=17, theta_labels=True)
-    fig, ax = plt.subplots(subplot_kw={"projection": "polar"})
-    text = ax.text(0, 1, "label", clip_on=True)
+def test_coord_radial_setup_ax_sets_pad_and_unclips_text():
+    data = pd.DataFrame({"x": [1], "y": [1], "label": ["label"]})
+    p = (
+        ggplot(data, aes("x", "y", label="label"))
+        + geom_text()
+        + coord_radial(theta_label_pad=17, theta_labels=True)
+    )
 
+    fig = p.draw()
     try:
-        coord.post_setup_ax(ax)
+        ax = fig.axes[0]
         assert ax.xaxis.get_tick_params()["pad"] == 17
-        assert not text.get_clip_on()
+        assert not ax.texts[0].get_clip_on()
     finally:
         plt.close(fig)