feat: individual positioning of texts of a colorbar

Author: has2k1

doc/changelog.qmd

@@ -13,6 +13,9 @@ title: Changelog
 - You can now pass a sequence of horizontal and vertical alignment (ha & va) values in
   element_text for colorbars.
 
+- You can now use a sequence to set the position of each text on colorbar, or have the
+  positions alternate left and right or bottom and top.
+
 ### Bug Fixes
 
 - Fixed [](:class:`~plotnine.geom_smooth`) / [](:class:`~plotnine.stat_smooth`) when using a linear model via "lm" with weights for the model to do a weighted regression. This bug did not affect the formula API of the linear model. ({{< issue 1005 >}})

plotnine/guides/guide.py

@@ -18,6 +18,7 @@
     from typing_extensions import Self
 
     from plotnine import aes, guides
+    from plotnine.iapi import guide_text
     from plotnine.layer import Layers, layer
     from plotnine.scales.scale import scale
     from plotnine.typing import (
@@ -139,6 +140,13 @@ def _resolved_position_justification(
             just = cast("tuple[float, float]", just)
             return (pos, just)
 
+    @property
+    def num_breaks(self) -> int:
+        """
+        Number of breaks
+        """
+        return len(self.key)
+
     def train(
         self, scale: scale, aesthetic: Optional[str] = None
     ) -> Self | None:
@@ -177,7 +185,7 @@ class GuideElements:
     guide: guide
 
     @cached_property
-    def text(self):
+    def text(self) -> guide_text:
         raise NotImplementedError
 
     def __post_init__(self):
@@ -214,16 +222,16 @@ def title(self):
         )
 
     @cached_property
-    def text_position(self) -> Side:
+    def text_positions(self) -> Sequence[Side]:
         raise NotImplementedError
 
     @cached_property
-    def _text_margin(self) -> float:
+    def _text_margin(self) -> Sequence[float]:
         _margin = self.theme.getp(
             (f"legend_text_{self.guide_kind}", "margin")
         ).pt
-        _loc = get_opposite_side(self.text_position)[0]
-        return getattr(_margin, _loc)
+        locs = (get_opposite_side(p)[0] for p in self.text_positions)
+        return [getattr(_margin, loc) for loc in locs]
 
     @cached_property
     def title_position(self) -> Side:
@@ -279,33 +287,3 @@ def is_horizontal(self) -> bool:
         Whether the guide is horizontal
         """
         return self.direction == "horizontal"
-
-    def has(self, n: int) -> Sequence[str]:
-        """
-        Horizontal alignments per legend text
-        """
-        ha = self.text.ha
-        if isinstance(ha, (list, tuple)):
-            if len(ha) != n:
-                raise ValueError(
-                    "If `ha` is a sequence, its length should match the "
-                    f"number of texts. ({len(ha)} != {n})"
-                )
-        else:
-            ha = (ha,) * n
-        return ha
-
-    def vas(self, n: int) -> Sequence[str]:
-        """
-        Vertical alignments per legend texts
-        """
-        va = self.text.va
-        if isinstance(va, (list, tuple)):
-            if len(va) != n:
-                raise ValueError(
-                    "If `va` is a sequence, its length should match the "
-                    f"number of texts. ({len(va)} != {n})"
-                )
-        else:
-            va = (va,) * n
-        return va

plotnine/guides/guide_colorbar.py

@@ -11,6 +11,8 @@
 import pandas as pd
 from mizani.bounds import rescale
 
+from plotnine.iapi import guide_text
+
 from .._utils import get_opposite_side
 from ..exceptions import PlotnineError, PlotnineWarning
 from ..mapping.aes import rename_aesthetics
@@ -417,26 +419,26 @@ def add_labels(
     """
     from matplotlib.text import Text
 
-    n = len(labels)
-    sep = elements.text.margin
+    seps = elements.text.margins
     texts: list[Text] = []
-    has = elements.has(n)
-    vas = elements.vas(n)
+    has = elements.text.has
+    vas = elements.text.vas
+    width = elements.key_width
 
     # The horizontal and vertical alignments are set in the theme
     # or dynamically calculates in GuideElements and added to the
     # themeable properties dict
     if elements.is_vertical:
-        if elements.text_position == "right":
-            xs = [elements.key_width + sep] * n
-        else:
-            xs = [-sep] * n
+        xs = [
+            width + sep if side == "right" else -sep
+            for side, sep in zip(elements.text_positions, seps)
+        ]
     else:
         xs = ys
-        if elements.text_position == "bottom":
-            ys = [-sep] * n
-        else:
-            ys = [elements.key_width + sep] * n
+        ys = [
+            -sep if side == "bottom" else width + sep
+            for side, sep in zip(elements.text_positions, seps)
+        ]
 
     for x, y, s, ha, va in zip(xs, ys, labels, has, vas):
         t = Text(x, y, s, ha=ha, va=va)
@@ -475,43 +477,52 @@ def text(self):
         ha = self.theme.getp(("legend_text_colorbar", "ha"))
         va = self.theme.getp(("legend_text_colorbar", "va"))
         is_blank = self.theme.T.is_blank("legend_text_colorbar")
+        n = self.guide.num_breaks
 
         # Default text alignment depends on the direction of the
         # colorbar
-        _loc = get_opposite_side(self.text_position)
+        centers = ("center",) * n
+        has = (ha,) * n if isinstance(ha, str) else ha
+        vas = (va,) * n if isinstance(va, str) else va
+        opposite_sides = [get_opposite_side(s) for s in self.text_positions]
         if self.is_vertical:
-            ha = ha or _loc
-            va = va or "center"
+            has = has or opposite_sides
+            vas = vas or centers
         else:
-            va = va or _loc
-            ha = ha or "center"
-
-        return NS(
-            margin=self._text_margin,
-            align=None,
+            vas = vas or opposite_sides
+            has = has or centers
+        return guide_text(
+            self._text_margin,
+            aligns=centers,
             fontsize=size,
-            ha=ha,
-            va=va,
+            has=has,  # pyright: ignore[reportArgumentType]
+            vas=vas,  # pyright: ignore[reportArgumentType]
             is_blank=is_blank,
         )
 
     @cached_property
-    def text_position(self) -> Side:
-        if not (position := self.theme.getp("legend_text_position")):
+    def text_positions(self) -> Sequence[Side]:
+        if not (user_position := self.theme.getp("legend_text_position")):
             position = "right" if self.is_vertical else "bottom"
+            return (position,) * self.guide.num_breaks
 
-        if self.is_vertical and position not in ("right", "left"):
-            msg = (
-                "The text position for a vertical legend must be "
-                "either left or right."
-            )
-            raise PlotnineError(msg)
-        elif self.is_horizontal and position not in ("bottom", "top"):
-            msg = (
-                "The text position for a horizonta legend must be "
-                "either top or bottom."
+        alternate = {"left-right", "right-left", "bottom-top", "top-bottom"}
+        if user_position in alternate:
+            tup = user_position.split("-")
+            return [tup[i % 2] for i in range(self.guide.num_breaks)]
+
+        position = cast("Side | Sequence[Side]", user_position)
+
+        if isinstance(position, str):
+            position = (position,) * self.guide.num_breaks
+
+        valid = {"right", "left"} if self.is_vertical else {"bottom", "top"}
+        if any(p for p in position if p not in valid):
+            raise PlotnineError(
+                "The text position for a horizontal legend must be "
+                f"either one of {valid!r}. I got {user_position!r}."
             )
-            raise PlotnineError(msg)
+
         return position
 
     @cached_property

plotnine/guides/guide_legend.py

@@ -5,20 +5,21 @@
 from dataclasses import dataclass, field
 from functools import cached_property
 from itertools import islice
-from types import SimpleNamespace as NS
 from typing import TYPE_CHECKING, cast
 from warnings import warn
 
 import numpy as np
 import pandas as pd
 
+from plotnine.iapi import guide_text
+
 from .._utils import remove_missing
 from ..exceptions import PlotnineError, PlotnineWarning
 from ..mapping.aes import rename_aesthetics
 from .guide import GuideElements, guide
 
 if TYPE_CHECKING:
-    from typing import Any, Optional
+    from typing import Any, Optional, Sequence
 
     from matplotlib.artist import Artist
     from matplotlib.offsetbox import PackerBase
@@ -209,7 +210,7 @@ def _calculate_rows_and_cols(
         self, elements: GuideElementsLegend
     ) -> tuple[int, int]:
         nrow, ncol = self.nrow, self.ncol
-        nbreak = len(self.key)
+        nbreak = self.num_breaks
 
         if nrow and ncol:
             if nrow * ncol < nbreak:
@@ -248,7 +249,7 @@ def draw(self):
 
         obverse = slice(0, None)
         reverse = slice(None, None, -1)
-        nbreak = len(self.key)
+        nbreak = self.num_breaks
         targets = self.theme.targets
         keys_order = reverse if self.reverse else obverse
         elements = self.elements
@@ -259,8 +260,12 @@ def draw(self):
         targets.legend_title = title_box._text  # type: ignore
 
         # labels
-        props = {"ha": elements.text.ha, "va": elements.text.va}
-        labels = [TextArea(s, textprops=props) for s in self.key["label"]]
+        has = elements.text.has
+        vas = elements.text.vas
+        labels = [
+            TextArea(s, textprops={"ha": ha, "va": va})
+            for s, ha, va in zip(self.key["label"], has, vas)
+        ]
         _texts = [l._text for l in labels]  # type: ignore
         targets.legend_text_legend = _texts
 
@@ -287,18 +292,28 @@ def draw(self):
             "bottom": (VPacker, reverse),
             "top": (VPacker, obverse),
         }
-        packer, slc = lookup[elements.text_position]
+
         if self.elements.text.is_blank:
             key_boxes = [d for d in drawings][keys_order]
         else:
+            packers, slices = [], []
+            for side in elements.text_positions:
+                tup = lookup[side]
+                packers.append(tup[0])
+                slices.append(tup[1])
+
+            seps = elements.text.margins
+            aligns = elements.text.aligns
             key_boxes = [
                 packer(
                     children=[l, d][slc],
-                    sep=elements.text.margin,
-                    align=elements.text.align,
+                    sep=sep,
+                    align=align,
                     pad=0,
                 )
-                for d, l in zip(drawings, labels)
+                for d, l, packer, slc, sep, align in zip(
+                    drawings, labels, packers, slices, seps, aligns
+                )
             ][keys_order]
 
         # Put the entries together in rows or columns
@@ -326,7 +341,7 @@ def draw(self):
                 break
 
         chunk_boxes: list[Artist] = [
-            packer_dim1(children=chunk, align="left", sep=sep1, pad=0)
+            packer_dim1(children=chunk, align="right", sep=sep1, pad=0)
             for chunk in chunks
         ]
 
@@ -364,26 +379,38 @@ def text(self):
         ha = self.theme.getp(("legend_text_legend", "ha"), "center")
         va = self.theme.getp(("legend_text_legend", "va"), "center")
         is_blank = self.theme.T.is_blank("legend_text_legend")
+        n = self.guide.num_breaks
 
         # The original ha & va values are used by the HPacker/VPacker
         # to align the TextArea with the DrawingArea.
         # We set ha & va to values that combine best with the aligning
         # for the text area.
-        align = va if self.text_position in {"left", "right"} else ha
-        return NS(
-            margin=self._text_margin,
-            align=align,
+        has = (ha,) * n if isinstance(ha, str) else ha
+        vas = (va,) * n if isinstance(va, str) else va
+        aligns = [
+            va if side in ("right", "left") else ha
+            for side, ha, va in zip(self.text_positions, has, vas)
+        ]
+        return guide_text(
+            margins=self._text_margin,
+            aligns=aligns,  # pyright: ignore[reportArgumentType]
             fontsize=size,
-            ha="center",
-            va="baseline",
+            has=("center",) * n,
+            vas=("baseline",) * n,
             is_blank=is_blank,
         )
 
     @cached_property
-    def text_position(self) -> Side:
-        if not (pos := self.theme.getp("legend_text_position")):
-            pos = "right"
-        return pos
+    def text_positions(self) -> Sequence[Side]:
+        if not (position := self.theme.getp("legend_text_position")):
+            return ("right",) * self.guide.num_breaks
+
+        position = cast("Side | Sequence[Side]", position)
+
+        if isinstance(position, str):
+            position = (position,) * self.guide.num_breaks
+
+        return position
 
     @cached_property
     def key_spacing_x(self) -> float:

plotnine/guides/guides.py

@@ -31,12 +31,12 @@
     from plotnine.scales.scale import scale
     from plotnine.scales.scales import Scales
     from plotnine.typing import (
+        Justification,
         LegendPosition,
         NoGuide,
         Orientation,
         ScaledAestheticsName,
         Side,
-        TextJustification,
     )
 
     LegendOrColorbar: TypeAlias = (
@@ -437,7 +437,7 @@ def _position_inside(self) -> LegendPosition:
         return ensure_xy_location(just)
 
     @cached_property
-    def box_just(self) -> TextJustification:
+    def box_just(self) -> Justification | Literal["baseline"]:
         if not (box_just := self.theme.getp("legend_box_just")):
             box_just = (
                 "left" if self.position in {"left", "right"} else "right"