better engine and mode checking code

topepo · topepo · commit 91fb33e9578e · 2021-07-15T20:00:42.000-04:00
diff --git a/NEWS.md b/NEWS.md
@@ -1,14 +1,21 @@
 # parsnip (development version)
 
+## Model Specification Changes
+
+* A model function (`gen_additive_mod()`) was added for generalized additive models. 
+
 * Each model now has a default engine that is used when the model is defined. The default for each model is listed in the help documents. This also adds functionality to declare an engine in the model specification function. `set_engine()` is still required if engine-specific arguments need to be added. (#513)
 
+* parsnip now checks for a valid engine for a given mode (#529)
+
 * The default engine for `multinom_reg()` was changed to `nnet`. 
 
+## Other Changes
+
 * The helper functions `.convert_form_to_xy_fit()`, `.convert_form_to_xy_new()`, `.convert_xy_to_form_fit()`, and  `.convert_xy_to_form_new()` for converting between formula and matrix interface are now exported for developer use (#508).
 
 * Fix bug in `augment()` when non-predictor, non-outcome variables are included in data (#510).
 
-* A model function (`gen_additive_mod()`) was added for generalized additive models. 
 
 # parsnip 0.1.6
 
diff --git a/R/aaa_models.R b/R/aaa_models.R
@@ -1,5 +1,7 @@
 # Initialize model environments
 
+all_modes <- c("classification", "regression", "censored regression")
+
 # ------------------------------------------------------------------------------
 
 ## Rules about model-related information
@@ -134,25 +136,90 @@ check_mode_val <- function(mode) {
 }
 
 
-stop_incompatible_mode <- function(spec_modes) {
+stop_incompatible_mode <- function(spec_modes, eng) {
+  if (is.null(eng)) {
+    msg <- glue::glue(
+      "Available modes are: ",
+      glue::glue_collapse(glue::glue("'{spec_modes}'"), sep = ", ")
+    )
+  } else {
+    msg <- glue::glue(
+      "Available modes for engine {eng} are: ",
+      glue::glue_collapse(glue::glue("'{spec_modes}'"), sep = ", ")
+    )
+  }
+
+  rlang::abort(msg)
+}
+
+stop_incompatible_engine <- function(spec_engs, mode) {
   msg <- glue::glue(
-    "Available modes are: ",
-    glue::glue_collapse(glue::glue("'{spec_modes}'"), sep = ", ")
+    "Available engines for mode {mode} are: ",
+    glue::glue_collapse(glue::glue("'{spec_engs}'"), sep = ", ")
   )
   rlang::abort(msg)
 }
 
-# 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"))
+# check if class and mode and engine are compatible
+check_spec_mode_engine_val <- function(cls, eng, mode) {
+  all_modes <- c("unknown", all_modes)
+  if (!(mode %in% all_modes)) {
+    rlang::abort(paste0("'", mode, "' is not a known mode."))
+  }
+
+  model_info <- rlang::env_get(get_model_env(), cls)
+
+  # Cases where the model definition is in parsnip but all of the engines
+  # are contained in a different package
+  if (nrow(model_info) == 0) {
+    return(invisible(NULL))
+  }
+
+  # ------------------------------------------------------------------------------
+  # First check engine against any mode
+
+  spec_engs <- model_info$engine
+  # engine is allowed to be NULL
+  if (!is.null(eng) && !(eng %in% spec_engs)) {
+    rlang::abort(
+      paste0(
+        "Engine '", eng, "' is not supported for `", cls, "()`. See ",
+        "`show_engines('", cls, "')`."
+      )
+    )
+  }
+
+  # ----------------------------------------------------------------------------
+  # Check modes based on model and engine
+
+  spec_modes <- model_info$mode
+  if (!is.null(eng)) {
+    spec_modes <- spec_modes[model_info$engine == eng]
+  }
+  spec_modes <- unique(c("unknown", spec_modes))
+
   if (is.null(mode) || length(mode) > 1) {
-    stop_incompatible_mode(spec_modes)
+    stop_incompatible_mode(spec_modes, eng)
   } else if (!(mode %in% spec_modes)) {
-    stop_incompatible_mode(spec_modes)
+    stop_incompatible_mode(spec_modes, eng)
   }
+
+  # ----------------------------------------------------------------------------
+  # Check engine based on model and model
+
+  # How check for compatibility with the chosen mode (if any)
+  if (!is.null(mode) && mode != "unknown") {
+    spec_engs <- spec_engs[model_info$mode == mode]
+  }
+  spec_engs <- unique(spec_engs)
+  if (!is.null(eng) && !(eng %in% spec_engs)) {
+    stop_incompatible_engine(spec_engs, mode)
+  }
+
   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'`).")
@@ -625,8 +692,7 @@ get_dependency <- function(model) {
 set_fit <- function(model, mode, eng, value) {
   check_model_exists(model)
   check_eng_val(eng)
-  check_mode_val(mode)
-  check_engine_val(eng)
+  check_spec_mode_engine_val(model, eng, mode)
   check_fit_info(value)
 
   current <- get_model_env()
@@ -692,8 +758,7 @@ get_fit <- function(model) {
 set_pred <- function(model, mode, eng, type, value) {
   check_model_exists(model)
   check_eng_val(eng)
-  check_mode_val(mode)
-  check_engine_val(eng)
+  check_spec_mode_engine_val(model, eng, mode)
   check_pred_info(value, type)
 
   current <- get_model_env()
diff --git a/R/arguments.R b/R/arguments.R
@@ -81,7 +81,7 @@ set_mode <- function(object, mode) {
     spec_modes <- rlang::env_get(get_model_env(), paste0(cls, "_modes"))
     stop_incompatible_mode(spec_modes)
   }
-  check_spec_mode_val(cls, mode)
+  check_spec_mode_engine_val(cls, object$engine, mode)
   object$mode <- mode
   object
 }
diff --git a/R/engines.R b/R/engines.R
@@ -10,25 +10,6 @@ possible_engines <- function(object, ...) {
   unique(engs$engine)
 }
 
-stop_incompatible_engine <- function(avail_eng) {
-  msg <- glue::glue(
-    "Available engines are: ",
-    glue::glue_collapse(glue::glue("'{avail_eng}'"), sep = ", ")
-  )
-  rlang::abort(msg)
-}
-
-check_engine <- function(object) {
-  avail_eng <- possible_engines(object)
-  eng <- object$engine
-  if (is.null(eng) || length(eng) > 1) {
-    stop_incompatible_engine(avail_eng)
-  } else if (!(eng %in% avail_eng)) {
-    stop_incompatible_engine(avail_eng)
-  }
-  object
-}
-
 # ------------------------------------------------------------------------------
 
 shhhh <- function(x)
@@ -99,7 +80,7 @@ set_engine <- function(object, engine, ...) {
     stop_incompatible_engine(avail_eng)
   }
   object$engine <- engine
-  object <- check_engine(object)
+  check_spec_mode_engine_val(class(object)[1], object$engine, object$mode)
 
   if (object$engine == "liquidSVM") {
     lifecycle::deprecate_soft(
diff --git a/R/misc.R b/R/misc.R
@@ -23,9 +23,6 @@ check_empty_ellipse <- function (...)  {
   terms
 }
 
-all_modes <- c("classification", "regression", "censored regression")
-
-
 deparserizer <- function(x, limit = options()$width - 10) {
   x <- deparse(x, width.cutoff = limit)
   x <- gsub("^    ", "", x)
@@ -192,7 +189,7 @@ update_dot_check <- function(...) {
 #' @rdname add_on_exports
 new_model_spec <- function(cls, args, eng_args, mode, method, engine) {
 
-  check_spec_mode_val(cls, mode)
+  check_spec_mode_engine_val(cls, engine, mode)
 
   out <- list(args = args, eng_args = eng_args,
               mode = mode, method = method, engine = engine)
diff --git a/R/translate.R b/R/translate.R
@@ -59,14 +59,15 @@ translate.default <- function(x, engine = x$engine, ...) {
   mod_name <- specific_model(x)
 
   x$engine <- engine
-  x <- check_engine(x)
-
   if (x$mode == "unknown") {
     rlang::abort("Model code depends on the mode; please specify one.")
   }
 
-  if (is.null(x$method))
+  check_spec_mode_engine_val(class(x)[1], x$engine, x$mode)
+
+  if (is.null(x$method)) {
     x$method <- get_model_spec(mod_name, x$mode, engine)
+  }
 
   arg_key <- get_args(mod_name, engine)
 
@@ -174,7 +175,7 @@ deharmonize <- function(args, key) {
 
 add_methods <- function(x, engine) {
   x$engine <- engine
-  x <- check_engine(x)
+  check_spec_mode_engine_val(class(x)[1], x$engine, x$mode)
   x$method <- get_model_spec(specific_model(x), x$mode, x$engine)
   x
 }
diff --git a/tests/testthat/test_args_and_modes.R b/tests/testthat/test_args_and_modes.R
@@ -49,6 +49,41 @@ test_that('pipe engine', {
 test_that("can't set a mode that isn't allowed by the model spec", {
   expect_error(
     set_mode(linear_reg(), "classification"),
-    "Available modes are:"
+    "Available modes"
   )
 })
+
+
+
+test_that("unavailable modes for an engine and vice-versa", {
+  expect_error(
+    decision_tree() %>%
+      set_mode("regression") %>%
+      set_engine("C5.0"),
+    "Available modes for engine C5"
+  )
+  expect_error(
+    decision_tree() %>%
+      set_engine("C5.0") %>%
+      set_mode("regression"),
+    "Available modes for engine C5"
+  )
+
+  expect_error(
+    decision_tree(engine = NULL) %>%
+      set_engine("C5.0") %>%
+      set_mode("regression"),
+    "Available modes for engine C5"
+  )
+
+  expect_error(
+    decision_tree(engine = NULL)%>%
+      set_mode("regression") %>%
+      set_engine("C5.0"),
+    "Available modes for engine C5"
+  )
+
+})
+
+
+

Original file line number	Diff line number	Diff line change
`@@ -81,7 +81,7 @@ set_mode <- function(object, mode) {`
`81`	`81`	`spec_modes <- rlang::env_get(get_model_env(), paste0(cls, "_modes"))`
`82`	`82`	`stop_incompatible_mode(spec_modes)`
`83`	`83`	`}`
`84`		`- check_spec_mode_val(cls, mode)`
	`84`	`+ check_spec_mode_engine_val(cls, object$engine, mode)`
`85`	`85`	`object$mode <- mode`
`86`	`86`	`object`
`87`	`87`	`}`