Train custom classes

NEWS.md

@@ -1,4 +1,7 @@
 # scales (development version)
+* Range training now supports custom classes. Continuous classes require a
+  `range()` method that returns numeric values. Discrete classes require a
+  `levels()` method (and optionally a `droplevels()` method) (#480).
 * New `label_glue()` labelling function for interpolated strings (#457).
 * `fullseq()` and by extension `breaks_width()` can now deal with unsorted 
   ranges (#435).

R/scale-continuous.R

@@ -41,6 +41,11 @@ train_continuous <- function(new, existing = NULL, call = caller_env()) {
     return(existing)
   }
 
+  new <- try_fetch(
+    range(new, na.rm = TRUE),
+    error = function(cnd) new
+  )
+
   if (is.factor(new) || !typeof(new) %in% c("integer", "double")) {
     example <- unique(new)
     example <- example[seq_len(pmin(length(example), 5))]

R/scale-discrete.R

@@ -15,7 +15,7 @@ dscale <- function(x, palette, na.value = NA) {
 }
 
 is.discrete <- function(x) {
-  is.factor(x) || is.character(x) || is.logical(x)
+  !is.null(levels(x)) || is.character(x) || is.logical(x)
 }
 
 #' Train (update) a discrete scale
@@ -90,8 +90,8 @@ discrete_range <- function(old, new, drop = FALSE, na.rm = FALSE, fct = NA) {
 clevels <- function(x, drop = FALSE, na.rm = FALSE) {
   if (is.null(x)) {
     character()
-  } else if (is.factor(x)) {
-    if (drop) x <- factor(x)
+  } else if (!is.null(levels(x))) {
+    if (drop) x <- droplevels(x)
 
     values <- levels(x)
     if (na.rm) {

tests/testthat/_snaps/scale-continuous.md

@@ -5,5 +5,5 @@
     Condition
       Error:
       ! Discrete value supplied to a continuous scale.
-      i Example values: "A", "B", "C", "D", and "E".
+      i Example values: "A" and "E".
 

tests/testthat/test-scale-continuous.R

@@ -34,3 +34,13 @@ test_that("train_continuous works with integer64", {
     c(1, 10)
   )
 })
+
+test_that("train_continuous can train on S3 classes", {
+  my_obj <- structure(c("IX", "CM", "X", "IV"), class = "bar")
+  range.bar <- function(x, ...) range(.romans[x], ...)
+  registerS3method("range", "bar", method = range.bar)
+  expect_equal(
+    train_continuous(my_obj),
+    c(4, 900)
+  )
+})

tests/testthat/test-scale-discrete.R

@@ -27,3 +27,13 @@ test_that("na.rm = TRUE drops NA", {
   expect_equal(train_discrete(x2, na.rm = TRUE), "a")
   expect_equal(train_discrete(x3, na.rm = TRUE), "a")
 })
+
+test_that("discrete ranges can be trained on S3 classes", {
+  my_obj <- structure(list("A", c("B", "C")), class = "foo")
+  levels.foo <- function(x) unique(unlist(x))
+  registerS3method("levels", "foo", method = levels.foo)
+  expect_equal(
+    train_discrete(my_obj),
+    LETTERS[1:3]
+  )
+})