merge pr #723: error on predict(type = "prob") with outcome level named "class"

Author: simonpcouch

NEWS.md

@@ -10,6 +10,8 @@
 
 *  An inconsistency for probability type predictions for two-class GAM models was fixed (#708)
 
+* `predict(type = "prob")` will now provide an error if the outcome variable has a level called `"class"` (#720).
+
 # parsnip 0.2.1
 
 * Fixed a major bug in spark models induced in the previous version (#671).

R/predict_classprob.R

@@ -9,7 +9,7 @@ predict_classprob.model_fit <- function(object, new_data, ...) {
     rlang::abort("`predict.model_fit()` is for predicting factor outcomes.")
 
   check_spec_pred_type(object, "prob")
-
+  check_spec_levels(object)
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")
@@ -48,3 +48,16 @@ predict_classprob.model_fit <- function(object, new_data, ...) {
 # @inheritParams predict.model_fit
 predict_classprob <- function(object, ...)
   UseMethod("predict_classprob")
+
+check_spec_levels <- function(spec) {
+  if ("class" %in% spec$lvl) {
+    rlang::abort(
+      glue::glue(
+        "The outcome variable `{spec$preproc$y_var}` has a level called 'class'. ",
+        "This value is reserved for parsnip's classification internals; please ",
+        "change the levels, perhaps with `forcats::fct_relevel()`."
+      ),
+      call = NULL
+    )
+  }
+}

tests/testthat/test_predict_formats.R

@@ -56,6 +56,31 @@ test_that('non-standard levels', {
                c("2low", "high+values"))
 })
 
+test_that('predict(type = "prob") with level "class" (see #720)', {
+  x <- tibble::tibble(
+    boop = factor(sample(c("class", "class_1"), 100, replace = TRUE)),
+    bop = rnorm(100),
+    beep = rnorm(100)
+  )
+
+  expect_error(
+    regexp = NA,
+    mod <- logistic_reg() %>%
+      set_mode(mode = "classification") %>%
+      fit(boop ~ bop + beep, data = x)
+  )
+
+  expect_error(
+    regexp = NA,
+    predict(mod, type = "class", new_data = x)
+  )
+
+  expect_error(
+    regexp = "variable `boop` has a level called 'class'",
+    predict(mod, type = "prob", new_data = x)
+  )
+})
+
 
 test_that('non-factor classification', {
   skip_if(run_glmnet)