Extend palette contract: Palette capabilities

R/ggproto.r

@@ -351,3 +351,7 @@ format.ggproto_method <- function(x, ...) {
 
 # proto2 TODO: better way of getting formals for self$draw
 ggproto_formals <- function(x) formals(environment(x)$f)
+
+ggproto_attr <- function(x, which, default = NULL) {
+  attr(environment(x)$f, which = which, exact = TRUE) %||% default
+}

R/scale-.r

@@ -609,7 +609,14 @@ ScaleContinuous <- ggproto("ScaleContinuous", Scale,
     pal <- self$palette(uniq)
     scaled <- pal[match(x, uniq)]
 
-    ifelse(!is.na(scaled), scaled, self$na.value)
+    # A specific palette can have as attribute "may_return_NA = FALSE"
+    # If it has such attribute, we will skip the ifelse(!is.na(scaled), ...)
+    pal_may_return_na <- ggproto_attr(self$palette, "may_return_NA", default = TRUE)
+    if (pal_may_return_na) {
+      scaled <- ifelse(!is.na(scaled), scaled, self$na.value)
+    }
+
+    scaled
   },
 
   rescale = function(self, x, limits = self$get_limits(), range = limits) {

tests/testthat/test-scale-colour-continuous.R

@@ -18,3 +18,32 @@ test_that("type argument is checked for proper input", {
     scale_colour_continuous(type = "abc")
   )
 })
+
+test_that("palette with may_return_NA=FALSE works as expected", {
+  sc <- scale_fill_continuous()
+  # A palette that may return NAs, will have NAs replaced by the scale's na.value
+  # by the scale:
+  sc$palette <- structure(
+    function(x) {
+      rep(NA_character_, length(x))
+    },
+    may_return_NA = TRUE
+  )
+  sc$na.value <- "red"
+  nat <- sc$map(0.5, limits = c(0, 1))
+  expect_equal(nat, "red")
+
+  # This palette is lying, because it returns NA even though it says it can't.
+  # The scale will not replace the NA values, leading to further errors.
+  # You should not do this in production, but it helps to test:
+  sc <- scale_fill_continuous()
+  sc$palette <- structure(
+    function(x) {
+      rep(NA_character_, length(x))
+    },
+    may_return_NA = FALSE
+  )
+  sc$na.value <- "red"
+  nat <- sc$map(0.5, limits = c(0, 1))
+  expect_equal(nat, NA_character_)
+})