Add helpers.get_aesthetic_limits

closes #906

Author: has2k1

doc/_quartodoc.yml

@@ -562,6 +562,13 @@ quartodoc:
         - figure_size
         - figure_format
 
+    - title: Tools
+      desc: |
+        Functions that you may occasioanally find helpful
+      package: plotnine.helpers
+      contents:
+          - get_aesthetic_limits
+
     - title: Datasets
       desc: |
         These datasets ship with the plotnine and you can import them with

doc/changelog.qmd

@@ -76,6 +76,10 @@ title: Changelog
 
 - Gained the Anscombe Quartet dataset, [](:attr:`~plotnine.data.anscombe_quartet`).
 
+- Gained new module `plotnine.helpers` to be a place for some helper functions.
+
+- Gained new helper function [](:func:`~plotnine.helpers.get_aesthetic_limits`).
+
 ### Enhancements
 
 - Included datasets [](:attr:`~plotnine.data.mpg`), [](:attr:`~plotnine.data.msleep`) and

plotnine/helpers.py

@@ -0,0 +1,49 @@
+from __future__ import annotations
+
+from copy import deepcopy
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    from typing import Sequence
+
+    from plotnine import ggplot
+
+__all__ = ("get_aesthetic_limits",)
+
+
+def get_aesthetic_limits(
+    plot: ggplot,
+    ae: str,
+) -> (
+    tuple[float, float]
+    | Sequence[str]
+    | list[tuple[float]]
+    | list[Sequence[str]]
+):
+    """
+    Get the limits of an aesthetic
+
+    These are the limits before they are expanded.
+
+    Parameters
+    ----------
+    plot :
+        ggplot object
+
+    ae :
+        Name of aesthetic
+
+    Returns
+    -------
+    out :
+        The limits of the aesthetic. If the plot is facetted, (has many
+        panels), it is a sequence of limits, one for each panel.
+    """
+    plot = deepcopy(plot)
+    plot._build()
+    limits = [
+        getattr(panel, ae).limits
+        for panel in plot._build_objs.layout.panel_params
+    ]
+
+    return limits[0] if len(limits) == 1 else limits

tests/test_helpers.py

@@ -0,0 +1,40 @@
+from typing import cast
+
+import numpy.testing as npt
+import pandas as pd
+
+from plotnine import aes, facet_wrap, geom_bar, geom_point, ggplot
+from plotnine.helpers import get_aesthetic_limits
+
+
+class TestGetAestheticLimits:
+    data = pd.DataFrame(
+        {
+            "x": [0, 1, 2, 3, 4, 5, 6],
+            "y": [0, 1, 2, 3, 4, 5, 6],
+            "g": list("aabbbcc"),
+        }
+    )
+
+    def test_continuous_limits(self):
+        p = ggplot(self.data, aes("x", "y")) + geom_point()
+        limits = cast("tuple[float, float]", get_aesthetic_limits(p, "x"))
+        npt.assert_array_almost_equal(limits, [0, 6])
+
+    def test_discrete_limits(self):
+        p = ggplot(self.data, aes("g")) + geom_bar()
+        limits = cast("list[str]", get_aesthetic_limits(p, "x"))
+        assert limits == ["a", "b", "c"]
+
+    def test_facet_limits(self):
+        p = (
+            ggplot(self.data, aes("x", "y"))
+            + geom_point()
+            + facet_wrap("g", scales="free_x")
+        )
+        limits = cast(
+            "list[tuple[float, float]]", get_aesthetic_limits(p, "x")
+        )
+        npt.assert_array_almost_equal(limits[0], [0, 1])
+        npt.assert_array_almost_equal(limits[1], [2, 4])
+        npt.assert_array_almost_equal(limits[2], [5, 6])