From bd62087926f3ac22b9444ce14152e1e083e89391 Mon Sep 17 00:00:00 2001
From: simonpcouch <simonpatrickcouch@gmail.com>
Date: Fri, 19 Jan 2024 15:53:19 -0600
Subject: [PATCH] improve error message for unsupported model modes

---
 NEWS.md                                 |  2 ++
 R/add_candidates.R                      |  8 ++++++++
 R/utils.R                               |  2 +-
 tests/testthat/_snaps/add_candidates.md |  8 ++++++++
 tests/testthat/test_add_candidates.R    | 10 ++++++++++
 5 files changed, 29 insertions(+), 1 deletion(-)

diff --git a/NEWS.md b/NEWS.md
index 5c15fb93..aca8d2ac 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,5 +1,7 @@
 # stacks (development version)
 
+* Improved error message for unsupported model modes (#152).
+
 # stacks 1.0.3
 
 * Refine package alt text (#202).
diff --git a/R/add_candidates.R b/R/add_candidates.R
index 9dce251a..91062e85 100644
--- a/R/add_candidates.R
+++ b/R/add_candidates.R
@@ -222,6 +222,14 @@ add_candidates.default <- function(data_stack, candidates, name, ...) {
   new_mode <- wf_spec$mode
   old_mode <- attr(stack, "mode")
   
+  if (isFALSE(new_mode %in% c("regression", "classification"))) {
+    cli_abort(
+      "The {.pkg stacks} package does not support stacking models with mode
+       {.val {new_mode}}.",
+      call = NULL
+    )
+  }
+  
   attr(stack, "mode") <- new_mode
   
   stack
diff --git a/R/utils.R b/R/utils.R
index 78337c06..02c44f13 100644
--- a/R/utils.R
+++ b/R/utils.R
@@ -121,7 +121,7 @@ should_run_examples <- function(suggests = NULL) {
 }
 
 mode_is_regression <- function(x) {
-  isTRUE(x[["mode"]] %in% c("regression", "censored regression"))
+  isTRUE(x[["mode"]] %in% "regression")
 }
 
 
diff --git a/tests/testthat/_snaps/add_candidates.md b/tests/testthat/_snaps/add_candidates.md
index 4535acb0..b172b061 100644
--- a/tests/testthat/_snaps/add_candidates.md
+++ b/tests/testthat/_snaps/add_candidates.md
@@ -115,6 +115,14 @@
       Error:
       ! The supplied candidates were tuned/fitted using only metrics that rely on hard class predictions. Please tune/fit with at least one class probability-based metric, such as `roc_auc` (`?yardstick::roc_auc()`).
 
+---
+
+    Code
+      stacks() %>% add_candidates(reg_res_lr)
+    Condition
+      Error:
+      ! The stacks package does not support stacking models with mode "censored regression".
+
 # model definition naming works as expected
 
     Code
diff --git a/tests/testthat/test_add_candidates.R b/tests/testthat/test_add_candidates.R
index 5abe9874..f8c890af 100644
--- a/tests/testthat/test_add_candidates.R
+++ b/tests/testthat/test_add_candidates.R
@@ -209,6 +209,16 @@ test_that("add_candidates errors informatively with bad arguments", {
     error = TRUE
   )
   
+  # fake a censored regression tuning result
+  wf <- attr(reg_res_lr, "workflow")
+  wf$fit$actions$model$spec$mode <- "censored regression"
+  attr(reg_res_lr, "workflow") <- wf
+  
+  expect_snapshot(
+    stacks() %>% add_candidates(reg_res_lr),
+    error = TRUE
+  )
+  
   # warn when stacking may fail due to tuning failure
   # TODO: re-implement tests--devel tune now errors on previous failure
 })