From 470c1d2db0d44e3e778588787606115dc3f232a2 Mon Sep 17 00:00:00 2001
From: Hassan Kibirige <has2k1@gmail.com>
Date: Mon, 6 May 2024 15:26:27 +0300
Subject: [PATCH] ENH: Fewer spurious legend + after_scale warnings

---
 doc/changelog.qmd               |  8 ++++++++
 plotnine/ggplot.py              |  2 +-
 plotnine/guides/guide_legend.py | 18 +++++++++++++++---
 plotnine/layer.py               | 18 ++++--------------
 tests/test_guide_internals.py   | 16 ++++++++++++++++
 5 files changed, 44 insertions(+), 18 deletions(-)
 create mode 100644 tests/test_guide_internals.py

diff --git a/doc/changelog.qmd b/doc/changelog.qmd
index f5ef3c866..3f62e23eb 100644
--- a/doc/changelog.qmd
+++ b/doc/changelog.qmd
@@ -1,6 +1,14 @@
 ---
 title: Changelog
 ---
+## v0.13.6
+(not-yet-released)
+
+### Enhancements
+
+- Stopped spurious warnings of the form ``PlotnineWarning: Failed to apply
+  `after_scale` modifications to the legend.`` when the `after_scale`
+  mapping is for another aesthetic.
 
 ## v0.13.5
 (2024-04-26)
diff --git a/plotnine/ggplot.py b/plotnine/ggplot.py
index a9c9ce11a..efd3e5353 100755
--- a/plotnine/ggplot.py
+++ b/plotnine/ggplot.py
@@ -400,7 +400,7 @@ def _build(self):
         layout.setup_panel_params(self.coordinates)
 
         # fill in the defaults
-        layers.use_defaults()
+        layers.use_defaults_after_scale()
 
         # Allow stats to modify the layer data
         layers.finish_statistics()
diff --git a/plotnine/guides/guide_legend.py b/plotnine/guides/guide_legend.py
index 995097b8d..90c52add8 100644
--- a/plotnine/guides/guide_legend.py
+++ b/plotnine/guides/guide_legend.py
@@ -160,23 +160,35 @@ def create_geoms(self):
                 continue
 
             matched = self.legend_aesthetics(l)
+            matched_set = set(matched)
 
             # This layer does not contribute to the legend
-            if not set(matched) - exclude:
+            if not matched_set - exclude:
                 continue
 
             data = self.key[matched].copy()
 
             # 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.
+            aes_modifiers = {
+                ae: expr
+                for ae, expr in l.mapping._scaled.items()
+                if ae in matched_set
+            }
+
             try:
-                data = l.use_defaults(data)
+                data = l.use_defaults(data, aes_modifiers)
             except PlotnineError:
                 warn(
                     "Failed to apply `after_scale` modifications "
                     "to the legend.",
                     PlotnineWarning,
                 )
-                data = l.use_defaults(data, aes_modifiers={})
+                data = l.use_defaults(data, {})
 
             # override.aes in guide_legend manually changes the geom
             for ae in set(self.override_aes) & set(data.columns):
diff --git a/plotnine/layer.py b/plotnine/layer.py
index 79ad1358d..9236061f9 100644
--- a/plotnine/layer.py
+++ b/plotnine/layer.py
@@ -368,8 +368,8 @@ def draw(self, layout: Layout, coord: coord):
 
     def use_defaults(
         self,
-        data: pd.DataFrame | None = None,
-        aes_modifiers: dict[str, Any] | None = None,
+        data: pd.DataFrame,
+        aes_modifiers: dict[str, Any],
     ) -> pd.DataFrame:
         """
         Prepare/modify data for plotting
@@ -382,12 +382,6 @@ def use_defaults(
             Expression to evaluate and replace aesthetics in
             the data.
         """
-        if data is None:
-            data = self.data
-
-        if aes_modifiers is None:
-            aes_modifiers = self.mapping._scaled
-
         return self.geom.use_defaults(data, aes_modifiers)
 
     def finish_statistics(self):
@@ -474,13 +468,9 @@ def compute_position(self, layout: Layout):
         for l in self:
             l.compute_position(layout)
 
-    def use_defaults(
-        self,
-        data: pd.DataFrame | None = None,
-        aes_modifiers: dict[str, Any] | None = None,
-    ):
+    def use_defaults_after_scale(self):
         for l in self:
-            l.use_defaults(data, aes_modifiers)
+            l.use_defaults(l.data, l.mapping._scaled)
 
     def transform(self, scales: Scales):
         for l in self:
diff --git a/tests/test_guide_internals.py b/tests/test_guide_internals.py
new file mode 100644
index 000000000..e43ae9fd6
--- /dev/null
+++ b/tests/test_guide_internals.py
@@ -0,0 +1,16 @@
+import warnings
+
+from plotnine import (
+    aes,
+    geom_point,
+    ggplot,
+)
+from plotnine.data import mtcars
+
+
+def test_no_after_scale_warning():
+    p = ggplot(mtcars, aes("wt", "mpg")) + geom_point()
+
+    with warnings.catch_warnings():
+        warnings.simplefilter("error")
+        p.draw_test()  # type: ignore