Check valid spec mode

NEWS.md

@@ -12,8 +12,11 @@
 
 * Re-licensed package from GPL-2 to MIT. See [consent from copyright holders here](https://github.com/tidymodels/parsnip/issues/462).
 
+* `set_mode()` now checks if `mode` is compatible with the model class, similar to `new_model_spec()` (@jtlandis, #467).
+
 * Re-organized model documentation for `update` methods (#479).
 
+
 # parsnip 0.1.5
 
 * An RStudio add-in is available that makes writing multiple `parsnip` model specifications to the source window. It can be accessed via the IDE addin menus or by calling `parsnip_addin()`.

R/aaa_models.R

@@ -132,6 +132,19 @@ check_mode_val <- function(mode) {
   invisible(NULL)
 }
 
+# check if class and mode are compatible
+check_spec_mode_val <- function(cls, mode) {
+  spec_modes <- rlang::env_get(get_model_env(), paste0(cls, "_modes"))
+  if (!(mode %in% spec_modes))
+    rlang::abort(
+      glue::glue(
+        "`mode` should be one of: ",
+        glue::glue_collapse(glue::glue("'{spec_modes}'"), sep = ", ")
+      )
+    )
+  invisible(NULL)
+}
+
 check_engine_val <- function(eng) {
   if (rlang::is_missing(eng) || length(eng) != 1 || !is.character(eng))
     rlang::abort("Please supply a character string for an engine (e.g. `'lm'`).")

R/arguments.R

@@ -87,6 +87,7 @@ set_mode <- function(object, mode) {
       )
     )
   }
+  check_spec_mode_val(class(object)[1], mode)
   object$mode <- mode
   object
 }

R/misc.R

@@ -191,14 +191,8 @@ update_dot_check <- function(...) {
 #' @keywords internal
 #' @rdname add_on_exports
 new_model_spec <- function(cls, args, eng_args, mode, method, engine) {
-  spec_modes <- rlang::env_get(get_model_env(), paste0(cls, "_modes"))
-  if (!(mode %in% spec_modes))
-    rlang::abort(
-      glue::glue(
-        "`mode` should be one of: ",
-        glue::glue_collapse(glue::glue("'{spec_modes}'"), sep = ", ")
-      )
-    )
+
+  check_spec_mode_val(cls, mode)
 
   out <- list(args = args, eng_args = eng_args,
               mode = mode, method = method, engine = engine)

tests/testthat/test_args_and_modes.R

@@ -45,3 +45,10 @@ test_that('pipe engine', {
   expect_error(rand_forest() %>% set_mode(2))
   expect_error(rand_forest() %>% set_mode("haberdashery"))
 })
+
+test_that("can't set a mode that isn't allowed by the model spec", {
+  expect_error(
+    set_mode(linear_reg(), "classification"),
+    "`mode` should be one of"
+  )
+})