after_*() length checks

NEWS.md

@@ -88,6 +88,8 @@
   to different position aesthetics of the same axis (@teunbrand, #5818).
 * In `stat_bin()`, the default `boundary` is now chosen to better adhere to 
   the `nbin` argument (@teunbrand, #5882, #5036)
+* `after_stat()` and `after_scale()` throw warnings when the computed aesthetics
+  are not of the correct length (#5901).
 
 # ggplot2 3.5.1
 

R/geom-.R

@@ -167,6 +167,9 @@ Geom <- ggproto("Geom",
       }
 
       names(modified_aes) <- names(rename_aes(modifiers))
+
+      modified_aes <- cleanup_mismatched_data(modified_aes, nrow(data), "after_scale")
+
       modified_aes <- data_frame0(!!!compact(modified_aes))
 
       data <- cunion(modified_aes, data)

R/guide-legend.R

@@ -238,17 +238,7 @@ GuideLegend <- ggproto(
         is_modified <- is_scaled_aes(aesthetics) | is_staged_aes(aesthetics)
         modifiers   <- aesthetics[is_modified]
 
-        data <- try_fetch(
-          layer$geom$use_defaults(params$key[matched_aes],
-                                  layer_params, modifiers),
-          error = function(cnd) {
-            cli::cli_warn(
-              "Failed to apply {.fn after_scale} modifications to legend",
-              parent = cnd
-            )
-            layer$geom$use_defaults(params$key[matched_aes], layer_params, list())
-          }
-        )
+        data <- layer$geom$use_defaults(params$key[matched_aes], layer_params, modifiers)
         data$.draw <- keep_key_data(params$key, df, matched_aes, layer$show.legend)
       } else {
         reps <- rep(1, nrow(params$key))

R/layer.R

@@ -412,6 +412,7 @@ Layer <- ggproto("Layer", NULL,
     if (self$stat$retransform) {
       stat_data <- plot$scales$transform_df(stat_data)
     }
+    stat_data <- cleanup_mismatched_data(stat_data, nrow(data), "after_stat")
 
     cunion(stat_data, data)
   },
@@ -498,3 +499,18 @@ set_draw_key <- function(geom, draw_key = NULL) {
   ggproto("", geom, draw_key = draw_key)
 }
 
+cleanup_mismatched_data <- function(data, n, fun) {
+  failed <- !lengths(data) %in% c(0, 1, n)
+  if (!any(failed)) {
+    return(data)
+  }
+
+  failed <- names(data)[failed]
+  cli::cli_warn(
+    "Failed to apply {.fn {fun}} for the following \\
+    aesthetic{?s}: {.field {failed}}."
+  )
+
+  data[failed] <- NULL
+  data
+}

tests/testthat/test-aes-calculated.R

@@ -70,6 +70,28 @@ test_that("staged aesthetics warn appropriately for duplicated names", {
   expect_snapshot_warning(ggplot_build(p))
 })
 
+test_that("calculated aesthetics throw warnings when lengths mismatch", {
+
+  df <- data.frame(x = 1:2)
+
+  p <- ggplot(df, aes(x, x))
+
+  expect_warning(
+    ggplot_build(
+      p + geom_point(aes(colour = after_stat(c("A", "B", "C"))))
+    ),
+    "Failed to apply"
+  )
+
+  expect_warning(
+    ggplot_build(
+      p + geom_point(aes(colour = after_scale(c("red", "green", "blue"))))
+    ),
+    "Failed to apply"
+  )
+
+})
+
 test_that("A deprecated warning is issued when stat(var) or ..var.. is used", {
   p1 <- ggplot(NULL, aes(stat(foo)))
   expect_snapshot_warning(b1 <- ggplot_build(p1))