Let stage() work without after_stat or after_scale

R/aes-evaluation.r

@@ -179,9 +179,20 @@ strip_stage <- function(expr) {
   if (is_call(uq_expr, c("after_stat", "after_scale"))) {
     uq_expr[[2]]
   } else if (is_call(uq_expr, "stage")) {
-    # Prefer stat mapping if present, otherwise original mapping (fallback to
-    # scale mapping) but there should always be two arguments to stage()
-    uq_expr$after_stat %||% uq_expr$start %||% (if (is.null(uq_expr$after_scale)) uq_expr[[3]]) %||% uq_expr[[2]]
+    # Prefer stat mapping if present
+    # i.e. `b` for `stage(a, after_stat = b)`
+    if (!is.null(uq_expr$after_stat))
+      return(uq_expr$after_stat)
+
+    # The case when after_stat is positionally provided.
+    # Be aware that we need to reject `stage(a, after_scale = c)`.
+    # i.e. `b` for `stage(a, b)`
+    if (is.null(uq_expr$after_scale) && length(uq_expr) >= 3)
+      return(uq_expr[[3]])
+
+    # The case when stage() has no `after_stat`
+    # i.e. `stage(a, after_scale = c)` or `stage(a)`
+    uq_expr[[2]]
   } else {
     expr
   }

tests/testthat/test-aes-calculated.r

@@ -69,3 +69,26 @@ test_that("staged aesthetics warn appropriately for duplicated names", {
   # One warning in building due to `stage()`/`after_scale()`
   expect_snapshot_warning(ggplot_build(p))
 })
+
+test_that("stage() works with only the start arg (#4873)", {
+  do_strip_stage <- function(...) {
+    args <- enexprs(...)
+    stage_call <- expr(stage(!!!args))
+    strip_stage(stage_call)
+  }
+
+  # the standard usage
+  expect_equal(do_strip_stage(x, after_stat = y), expr(y))
+
+  # positional
+  expect_equal(do_strip_stage(x, y), expr(y))
+
+  # when with only the start arg, simply strips stage()
+  expect_equal(do_strip_stage(x), expr(x))
+
+  # when after_scale is specified without after_stat, return the start arg
+  expect_equal(do_strip_stage(x, after_scale = y), expr(x))
+
+  # when an explicit NULL is specified, use it
+  expect_equal(do_strip_stage(x, after_stat = NULL), expr(NULL))
+})