Mask stage() at the expression level

R/aes-evaluation.R

@@ -359,3 +359,39 @@ make_labels <- function(mapping) {
   }
   Map(default_label, names(mapping), mapping)
 }
+
+eval_aesthetics <- function(aesthetics, data, mask = NULL) {
+
+  env <- child_env(base_env())
+
+  # Here we mask functions, often to replace `stage()` with context appropriate
+  # functions `stage_calculated()`/`stage_scaled()`.
+  if (length(mask) > 0) {
+    aesthetics <- substitute_aes(aesthetics, mask_function, mask = mask)
+  }
+
+  evaled <- lapply(aesthetics, eval_tidy, data = data, env = env)
+  names(evaled) <- names(aesthetics)
+  compact(rename_aes(evaled))
+}
+
+# `mask` is a list of functions where `names(mask)` indicate names of functions
+# that need to be replaced, and `mask[[i]]` is the function to replace it
+# with.
+mask_function <- function(x, mask) {
+  if (!is.call(x)) {
+    return(x)
+  }
+  nms <- names(mask)
+  x[-1] <- lapply(x[-1], mask_function, mask = mask)
+  if (!is_call(x, nms)) {
+    return(x)
+  }
+  for (nm in nms) {
+    if (is_call(x, nm)) {
+      x[[1]] <- mask[[nm]]
+      return(x)
+    }
+  }
+}
+

R/aes.R

@@ -189,9 +189,9 @@ rename_aes <- function(x) {
   }
   x
 }
-substitute_aes <- function(x) {
+substitute_aes <- function(x, fun = standardise_aes_symbols, ...) {
   x <- lapply(x, function(aesthetic) {
-    as_quosure(standardise_aes_symbols(quo_get_expr(aesthetic)), env = environment(aesthetic))
+    as_quosure(fun(quo_get_expr(aesthetic), ...), env = environment(aesthetic))
   })
   class(x) <- "uneval"
   x

R/geom-.R

@@ -157,12 +157,10 @@ Geom <- ggproto("Geom",
     # This order means that they will have access to all default aesthetics
     if (length(modifiers) != 0) {
       # Set up evaluation environment
-      env <- child_env(baseenv(), after_scale = after_scale)
-      # Mask stage with stage_scaled so it returns the correct expression
-      stage_mask <- child_env(emptyenv(), stage = stage_scaled)
-      mask <- new_data_mask(as_environment(data, stage_mask), stage_mask)
-      mask$.data <- as_data_pronoun(mask)
-      modified_aes <- lapply(substitute_aes(modifiers),  eval_tidy, mask, env)
+      modified_aes <- eval_aesthetics(
+        substitute_aes(modifiers), data,
+        mask = list(stage = stage_scaled)
+      )
 
       # Check that all output are valid data
       nondata_modified <- check_nondata_cols(modified_aes)
@@ -177,11 +175,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))
+      modified_aes <- data_frame0(!!!modified_aes)
 
       data <- cunion(modified_aes, data)
     }

R/layer.R

@@ -58,8 +58,8 @@
 #'   `NA`, the default, includes if any aesthetics are mapped.
 #'   `FALSE` never includes, and `TRUE` always includes.
 #'   It can also be a named logical vector to finely select the aesthetics to
-#'   display. To include legend keys for all levels, even 
-#'   when no data exists, use `TRUE`.  If `NA`, all levels are shown in legend, 
+#'   display. To include legend keys for all levels, even
+#'   when no data exists, use `TRUE`.  If `NA`, all levels are shown in legend,
 #'   but unobserved levels are omitted.
 #' @param inherit.aes If `FALSE`, overrides the default aesthetics,
 #'   rather than combining with them. This is most useful for helper functions
@@ -303,10 +303,7 @@ Layer <- ggproto("Layer", NULL,
     }
 
     # Evaluate aesthetics
-    env <- child_env(baseenv(), stage = stage)
-    evaled <- lapply(aesthetics, eval_tidy, data = data, env = env)
-    evaled <- compact(evaled)
-
+    evaled <- eval_aesthetics(aesthetics, data)
     plot$scales$add_defaults(evaled, plot$plot_env)
 
     # Check for discouraged usage in mapping
@@ -386,14 +383,10 @@ Layer <- ggproto("Layer", NULL,
     data_orig <- plot$scales$backtransform_df(data)
 
     # Add map stat output to aesthetics
-    env <- child_env(baseenv(), stat = stat, after_stat = after_stat)
-    stage_mask <- child_env(emptyenv(), stage = stage_calculated)
-    mask <- new_data_mask(as_environment(data_orig, stage_mask), stage_mask)
-    mask$.data <- as_data_pronoun(mask)
-
-    new <- substitute_aes(new)
-    stat_data <- lapply(new, eval_tidy, mask, env)
-
+    stat_data <- eval_aesthetics(
+      substitute_aes(new), data_orig,
+      mask = list(stage = stage_calculated)
+    )
     # Check that all columns in aesthetic stats are valid data
     nondata_stat_cols <- check_nondata_cols(stat_data)
     if (length(nondata_stat_cols) > 0) {
@@ -407,8 +400,7 @@ Layer <- ggproto("Layer", NULL,
       ))
     }
 
-    names(stat_data) <- names(new)
-    stat_data <- data_frame0(!!!compact(stat_data))
+    stat_data <- data_frame0(!!!stat_data)
 
     # Add any new scales, if needed
     plot$scales$add_defaults(stat_data, plot$plot_env)