Fix after_scale aesthetic to show up in the legend

fixes #926

Author: has2k1

doc/changelog.qmd

@@ -87,6 +87,10 @@ title: Changelog
 - Fixed bug where facetting along a column with unused categories could raise
   an error. {{< issue 930 >}}
 
+- Fixed bug where computed aesthetic from `after_scale` are were not applied to the legend.
+  {{< issue 926 >}}
+
+
 ## v0.14.5
 (2025-01-02)
 [![](https://zenodo.org/badge/DOI/10.5281/zenodo.14587381.svg)](https://doi.org/10.5281/zenodo.14587381)

plotnine/guides/guide.py

@@ -18,7 +18,7 @@
     from typing_extensions import Self
 
     from plotnine import aes, guides
-    from plotnine.layer import Layers
+    from plotnine.layer import Layers, layer
     from plotnine.scales.scale import scale
     from plotnine.typing import (
         LegendPosition,
@@ -79,7 +79,7 @@ def __post_init__(self):
         self.elements = cast("GuideElements", None)
         self.guides_elements: GuidesElements
 
-    def legend_aesthetics(self, layer):
+    def legend_aesthetics(self, layer: layer):
         """
         Return the aesthetics that contribute to the legend
 

plotnine/guides/guide_legend.py

@@ -171,21 +171,20 @@ def create_geoms(self):
             # Modify aesthetics
 
             # When doing after_scale evaluations, we only consider those
-            # for the aesthetics of this legend. The reduces the spurious
-            # warnings where an evaluation of another aesthetic failed yet
-            # it is not needed.
+            # for the aesthetics that are valid for this layer/geom.
             aes_modifiers = {
-                ae: expr
-                for ae, expr in l.mapping._scaled.items()
-                if ae in matched_set
+                ae: l.mapping._scaled[ae]
+                for ae in l.geom.aesthetics() & l.mapping._scaled.keys()
             }
 
             try:
                 data = l.use_defaults(data, aes_modifiers)
             except PlotnineError:
                 warn(
-                    "Failed to apply `after_scale` modifications "
-                    "to the legend.",
+                    "Failed to apply `after_scale` modifications to the "
+                    "legend. This probably should not happen. Help us "
+                    "discover why, please open and issue at "
+                    "https://github.com/has2k1/plotnine/issues",
                     PlotnineWarning,
                 )
                 data = l.use_defaults(data, {})

plotnine/mapping/aes.py

@@ -6,6 +6,7 @@
 from contextlib import suppress
 from copy import deepcopy
 from dataclasses import fields
+from functools import cached_property
 from typing import Any, Dict
 
 import pandas as pd
@@ -237,7 +238,7 @@ def _convert_deprecated_expr(self, kwargs):
                 kwargs[name] = after_stat(_after_stat)
         return kwargs
 
-    @property
+    @cached_property
     def _starting(self) -> dict[str, Any]:
         """
         Return the subset of aesthetics mapped from the layer data
@@ -254,7 +255,7 @@ def _starting(self) -> dict[str, Any]:
 
         return d
 
-    @property
+    @cached_property
     def _calculated(self) -> dict[str, Any]:
         """
         Return only the aesthetics mapped to calculated statistics
@@ -269,7 +270,7 @@ def _calculated(self) -> dict[str, Any]:
 
         return d
 
-    @property
+    @cached_property
     def _scaled(self) -> dict[str, Any]:
         """
         Return only the aesthetics mapped to after scaling

tests/baseline_images/test_guide_internals/guide_legend_after_scale.png

@@ -1,7 +1,11 @@
 import warnings
 
+import pandas as pd
+
 from plotnine import (
     aes,
+    after_scale,
+    geom_bar,
     geom_point,
     ggplot,
 )
@@ -14,3 +18,26 @@ def test_no_after_scale_warning():
     with warnings.catch_warnings():
         warnings.simplefilter("error")
         p.draw_test()  # type: ignore
+
+
+def test_guide_legend_after_scale():
+    def alphen(series, a):
+        ha = f"{round(a * 255):#04X}"[2:]
+        return [f"{hex_color}{ha}" for hex_color in series]
+
+    data = pd.DataFrame(
+        {"var1": [1, 2, 2, 3, 3, 3, 4, 4, 4, 4, 5, 5, 5, 5, 5]}
+    )
+    p = (
+        ggplot(
+            data,
+            aes(
+                "var1",
+                color="factor(var1)",
+                fill=after_scale("alphen(color, .5)"),
+            ),
+        )
+        + geom_bar()
+    )
+
+    assert p == "guide_legend_after_scale"