From fda519e2ac0828048c24f3c2b254c2607f75ac5c Mon Sep 17 00:00:00 2001
From: Teun van den Brand <49372158+teunbrand@users.noreply.github.com>
Date: Tue, 22 Oct 2024 10:36:50 +0200
Subject: [PATCH] Palette manipulation functions (#466)

---
 NAMESPACE                          | 10 +++++
 R/colour-manip.R                   | 59 ++++++++++++++++++++++++++++--
 R/pal-.R                           | 16 +++++++-
 man/col_mix.Rd                     |  4 +-
 man/colour_manip.Rd                |  2 +-
 tests/testthat/test-colour-manip.R | 11 ++++++
 6 files changed, 94 insertions(+), 8 deletions(-)

diff --git a/NAMESPACE b/NAMESPACE
index 64ddd138..a3328b77 100644
--- a/NAMESPACE
+++ b/NAMESPACE
@@ -8,6 +8,14 @@ S3method(as_discrete_pal,"function")
 S3method(as_discrete_pal,character)
 S3method(as_discrete_pal,default)
 S3method(as_discrete_pal,pal_continuous)
+S3method(col_lighter,default)
+S3method(col_lighter,scales_pal)
+S3method(col_mix,default)
+S3method(col_mix,scales_pal)
+S3method(col_saturate,default)
+S3method(col_saturate,scales_pal)
+S3method(col_shift,default)
+S3method(col_shift,scales_pal)
 S3method(fullseq,Date)
 S3method(fullseq,POSIXt)
 S3method(fullseq,difftime)
@@ -17,6 +25,8 @@ S3method(offset_by,Date)
 S3method(offset_by,POSIXt)
 S3method(offset_by,difftime)
 S3method(offset_by,numeric)
+S3method(plot,pal_continuous)
+S3method(plot,pal_discrete)
 S3method(plot,transform)
 S3method(print,transform)
 S3method(rescale,"NULL")
diff --git a/R/colour-manip.R b/R/colour-manip.R
index c4f15d85..95f8fdc6 100644
--- a/R/colour-manip.R
+++ b/R/colour-manip.R
@@ -121,7 +121,8 @@ show_col <- function(colours, labels = TRUE, borders = NULL, cex_label = 1,
 #'
 #' Produces an interpolation of two colours.
 #'
-#' @param a,b A character vector of colours.
+#' @param a Either a character vector of colours or a colour palette function.
+#' @param b A character vector of colours.
 #' @param amount A numeric fraction between 0 and 1 giving the contribution of
 #'   the `b` colour.
 #' @param space A string giving a colour space to perform mixing operation in.
@@ -139,6 +140,11 @@ show_col <- function(colours, labels = TRUE, borders = NULL, cex_label = 1,
 #' # Not recommended:
 #' col_mix("blue", "red", space = "hcl") # green!
 col_mix <- function(a, b, amount = 0.5, space = "rgb") {
+  UseMethod("col_mix")
+}
+
+#' @export
+col_mix.default <- function(a, b, amount = 0.5, space = "rgb") {
   input <- recycle_common(a = a, b = b, amount = amount)
   if (any(input$amount < 0 | input$amount > 1)) {
     cli::cli_abort("{.arg amount} must be between (0, 1).")
@@ -150,12 +156,17 @@ col_mix <- function(a, b, amount = 0.5, space = "rgb") {
   farver::encode_colour(new, alpha = alpha, from = space)
 }
 
+#' @export
+col_mix.scales_pal <- function(a, b, amount = 0.5, space = "rgb") {
+  wrap_col_adjustment(a, col_mix, list(b = b, amount = amount, space = space))
+}
+
 #' Colour manipulation
 #'
 #' These are a set of convenience functions for standard colour manipulation
 #' operations.
 #'
-#' @param col A character vector of colours.
+#' @param col A character vector of colours or a colour palette function.
 #' @param amount A numeric vector giving the change. The interpretation depends
 #'   on the function:
 #'   * `col_shift()` takes a number between -360 and 360 for shifting hues in
@@ -183,29 +194,69 @@ NULL
 #' @export
 #' @rdname colour_manip
 col_shift <- function(col, amount = 10) {
+  UseMethod("col_shift")
+}
+
+#' @export
+col_shift.default <- function(col, amount = 10) {
   input  <- recycle_common(col = col, amount = amount)
   new <- farver::decode_colour(input$col, alpha = TRUE, to = "hcl")
   new[, "h"] <- (new[, "h"] + input$amount) %% 360
   farver::encode_colour(new, new[, "alpha"], from = "hcl")
 }
 
+#' @export
+col_shift.scales_pal <- function(col, amount = 10) {
+  wrap_col_adjustment(col, col_shift, list(amount = amount))
+}
+
 #' @export
 #' @rdname colour_manip
 col_lighter <- function(col, amount = 10) {
+  UseMethod("col_lighter")
+}
+
+#' @export
+col_lighter.default <- function(col, amount = 10) {
   input <- recycle_common(col = col, amount = amount)
   farver::add_to_channel(input$col, "l", input$amount, space = "hsl")
 }
 
+#' @export
+col_lighter.scales_pal <- function(col, amount = 10) {
+  wrap_col_adjustment(col, col_lighter, list(amount = amount))
+}
+
 #' @export
 #' @rdname colour_manip
 col_darker <- function(col, amount = 10) {
-  input <- recycle_common(col = col, amount = amount)
-  farver::add_to_channel(input$col, "l", -input$amount, space = "hsl")
+  col_lighter(col, amount = -amount)
 }
 
 #' @export
 #' @rdname colour_manip
 col_saturate <- function(col, amount = 10) {
+  UseMethod("col_saturate")
+}
+
+#' @export
+col_saturate.default <- function(col, amount = 10) {
   input <- recycle_common(col = col, amount = amount)
   farver::add_to_channel(input$col, "s", input$amount, space = "hsl")
 }
+
+#' @export
+col_saturate.scales_pal <- function(col, amount = 10) {
+  wrap_col_adjustment(col, col_saturate, list(amount = amount))
+}
+
+wrap_col_adjustment <- function(inner, outer, args, call = caller_env()) {
+  check_object(inner, is_colour_pal, "a {.field colour} palette")
+  force_all(inner, outer, args)
+  fun <- function(...) inject(outer(inner(...), !!!args))
+  if (is_discrete_pal(inner)) {
+    new_discrete_palette(fun, type = "colour", nlevels = palette_nlevels(inner))
+  } else {
+    new_continuous_palette(fun, type = "colour", na_safe = palette_na_safe(inner))
+  }
+}
diff --git a/R/pal-.R b/R/pal-.R
index 4c20d174..66d1f27f 100644
--- a/R/pal-.R
+++ b/R/pal-.R
@@ -63,7 +63,7 @@
 #' show_col(pal(9))
 new_continuous_palette <- function(fun, type, na_safe = NA) {
   check_function(fun)
-  class(fun) <- union("pal_continuous", class(fun))
+  class(fun) <- union(c("pal_continuous", "scales_pal"), class(fun))
   attr(fun, "type") <- type
   attr(fun, "na_safe") <- na_safe
   fun
@@ -73,7 +73,7 @@ new_continuous_palette <- function(fun, type, na_safe = NA) {
 #' @export
 new_discrete_palette <- function(fun, type, nlevels = NA) {
   check_function(fun)
-  class(fun) <- union("pal_discrete", class(fun))
+  class(fun) <- union(c("pal_discrete", "scales_pal"), class(fun))
   attr(fun, "type") <- type
   attr(fun, "nlevels") <- nlevels
   fun
@@ -209,3 +209,15 @@ as_continuous_pal.character <- function(x, ...) {
   }
   as_continuous_pal(get_palette(x, ...))
 }
+
+# Utility -----------------------------------------------------------------
+
+#' @export
+plot.pal_discrete <- function(x, y, ..., n_max = 25) {
+  show_col(x(pmin(n_max, palette_nlevels(x))), ...)
+}
+
+#' @export
+plot.pal_continuous <- function(x, y, ..., n_max = 25) {
+  show_col(x(seq(0, 1, length.out = n_max)), ...)
+}
diff --git a/man/col_mix.Rd b/man/col_mix.Rd
index ae40f3ac..f5181402 100644
--- a/man/col_mix.Rd
+++ b/man/col_mix.Rd
@@ -7,7 +7,9 @@
 col_mix(a, b, amount = 0.5, space = "rgb")
 }
 \arguments{
-\item{a, b}{A character vector of colours.}
+\item{a}{Either a character vector of colours or a colour palette function.}
+
+\item{b}{A character vector of colours.}
 
 \item{amount}{A numeric fraction between 0 and 1 giving the contribution of
 the \code{b} colour.}
diff --git a/man/colour_manip.Rd b/man/colour_manip.Rd
index 2dae5274..ff023c69 100644
--- a/man/colour_manip.Rd
+++ b/man/colour_manip.Rd
@@ -17,7 +17,7 @@ col_darker(col, amount = 10)
 col_saturate(col, amount = 10)
 }
 \arguments{
-\item{col}{A character vector of colours.}
+\item{col}{A character vector of colours or a colour palette function.}
 
 \item{amount}{A numeric vector giving the change. The interpretation depends
 on the function:
diff --git a/tests/testthat/test-colour-manip.R b/tests/testthat/test-colour-manip.R
index a2f93225..744abb7e 100644
--- a/tests/testthat/test-colour-manip.R
+++ b/tests/testthat/test-colour-manip.R
@@ -72,3 +72,14 @@ test_that("col_saturate can (de)saturate colours", {
   expect_equal(col_saturate(x, -30), c("#996666", "#669966", "#666699"))
 })
 
+test_that("colour manipulation functions work on palettes", {
+
+  pal <- pal_manual(c("#FF0000", "#00FF00", "#0000FF"))
+
+  expect_equal(col_shift(pal, 180)(3),    c("#00B8B8", "#FF92FF", "#535300"))
+  expect_equal(col_darker(pal, 30)(3),    c("#660000", "#006600", "#000066"))
+  expect_equal(col_lighter(pal, 30)(3),   c("#FF9999", "#99FF99", "#9999FF"))
+  expect_equal(col_saturate(pal, -50)(3), c("#BF4040", "#40BF40", "#4040BF"))
+  expect_equal(col_mix(pal, "white")(3),  c("#FF8080", "#80FF80", "#8080FF"))
+
+})