modified checking method for correct pred type

Author: topepo

R/aaa_models.R

@@ -286,9 +286,13 @@ check_pred_info <- function(pred_obj, type) {
   invisible(NULL)
 }
 
-check_spec_pred_type <- function(object, type) {
+spec_has_pred_type <- function(object, type) {
   possible_preds <- names(object$spec$method$pred)
-  if (!any(possible_preds == type)) {
+  any(possible_preds == type)
+}
+check_spec_pred_type <- function(object, type) {
+  if (!spec_has_pred_type(object, type)) {
+    possible_preds <- names(object$spec$method$pred)
     rlang::abort(c(
       glue::glue("No {type} prediction method available for this model."),
       glue::glue("Value for `type` should be one of: ",

R/augment.R

@@ -6,8 +6,9 @@
 #' [fit()] and `new_data` contains the outcome column, a `.resid` column is
 #' also added.
 #'
-#' For classification models, the results include a column called `.pred_class`
-#' as well as class probability columns named `.pred_{level}`.
+#' For classification models, the results can include a column called
+#'  `.pred_class` as well as class probability columns named `.pred_{level}`.
+#'  This depends on what type of prediction types are available for the model.
 #' @param x A `model_fit` object produced by [fit()] or [fit_xy()].
 #' @param new_data A data frame or matrix.
 #' @param ... Not currently used.
@@ -56,6 +57,7 @@
 #'
 augment.model_fit <- function(x, new_data, ...) {
   if (x$spec$mode == "regression") {
+    check_spec_pred_type(x, "numeric")
     new_data <-
       new_data %>%
       dplyr::bind_cols(
@@ -68,13 +70,13 @@ augment.model_fit <- function(x, new_data, ...) {
       }
     }
   } else if (x$spec$mode == "classification") {
-    if (has_class_preds(x)) {
+    if (spec_has_pred_type(x, "class")) {
       new_data <- dplyr::bind_cols(
         new_data,
         predict(x, new_data = new_data, type = "class")
       )
     }
-    if (has_class_probs(x)) {
+    if (spec_has_pred_type(x, "prob")) {
       new_data <- dplyr::bind_cols(
         new_data,
         predict(x, new_data = new_data, type = "prob")
@@ -85,11 +87,3 @@ augment.model_fit <- function(x, new_data, ...) {
   }
   as_tibble(new_data)
 }
-
-has_class_preds <- function(x) {
-  any(names(x$spec$method$pred) == "class")
-}
-
-has_class_probs <- function(x) {
-  any(names(x$spec$method$pred) == "prob")
-}