Check the pred type in unified way across types

Author: juliasilge

R/predict_class.R

@@ -12,8 +12,7 @@ predict_class.model_fit <- function(object, new_data, ...) {
   if (object$spec$mode != "classification")
     rlang::abort("`predict.model_fit()` is for predicting factor outcomes.")
 
-  if (!any(names(object$spec$method$pred) == "class"))
-    rlang::abort("No class prediction module defined for this model.")
+  check_spec_pred_type(object, "class")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_classprob.R

@@ -9,8 +9,8 @@ predict_classprob.model_fit <- function(object, new_data, ...) {
   if (object$spec$mode != "classification")
     rlang::abort("`predict.model_fit()` is for predicting factor outcomes.")
 
-  if (!any(names(object$spec$method$pred) == "prob"))
-    rlang::abort("No class probability module defined for this model.")
+  check_spec_pred_type(object, "prob")
+
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_hazard.R

@@ -7,8 +7,7 @@
 predict_hazard.model_fit <-
   function(object, new_data, .time, ...) {
 
-    if (is.null(object$spec$method$pred$hazard))
-      rlang::abort("No hazard prediction method defined for this engine.")
+    check_spec_pred_type(object, "hazard")
 
     if (inherits(object$fit, "try-error")) {
       rlang::warn("Model fit failed; cannot make predictions.")

R/predict_interval.R

@@ -10,8 +10,7 @@
 #' @export
 predict_confint.model_fit <- function(object, new_data, level = 0.95, std_error = FALSE, ...) {
 
-  if (is.null(object$spec$method$pred$conf_int))
-    rlang::abort("No confidence interval method defined for this engine.")
+  check_spec_pred_type(object, "conf_int")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")
@@ -58,8 +57,7 @@ predict_confint <- function(object, ...)
 # @export
 predict_predint.model_fit <- function(object, new_data, level = 0.95, std_error = FALSE, ...) {
 
-  if (is.null(object$spec$method$pred$pred_int))
-    rlang::abort("No prediction interval method defined for this engine.")
+  check_spec_pred_type(object, "pred_int")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_linear_pred.R

@@ -6,8 +6,7 @@
 #' @export
 predict_linear_pred.model_fit <- function(object, new_data, ...) {
 
-  if (!any(names(object$spec$method$pred) == "linear_pred"))
-    rlang::abort("No prediction module defined for this model.")
+  check_spec_pred_type(object, "linear_pred")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_numeric.R

@@ -10,8 +10,7 @@ predict_numeric.model_fit <- function(object, new_data, ...) {
                             "Use `predict_class()` or `predict_classprob()` for ",
                             "classification models."))
 
-  if (!any(names(object$spec$method$pred) == "numeric"))
-    rlang::abort("No prediction module defined for this model.")
+  check_spec_pred_type(object, "numeric")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_quantile.R

@@ -9,8 +9,7 @@
 predict_quantile.model_fit <-
   function(object, new_data, quantile = (1:9)/10, ...) {
 
-    if (is.null(object$spec$method$pred$quantile))
-      rlang::abort("No quantile prediction method defined for this engine.")
+    check_spec_pred_type(object, "quantile")
 
     if (inherits(object$fit, "try-error")) {
       rlang::warn("Model fit failed; cannot make predictions.")

R/predict_raw.R

@@ -13,8 +13,7 @@ predict_raw.model_fit <- function(object, new_data, opts = list(), ...) {
       c(object$spec$method$pred$raw$args, opts)
   }
 
-  if (!any(names(object$spec$method$pred) == "raw"))
-    rlang::abort("No raw prediction module defined for this model.")
+  check_spec_pred_type(object, "raw")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")

R/predict_survival.R

@@ -7,8 +7,7 @@
 predict_survival.model_fit <-
   function(object, new_data, .time, ...) {
 
-    if (is.null(object$spec$method$pred$survival))
-      rlang::abort("No survival prediction method defined for this engine.")
+    check_spec_pred_type(object, "survival")
 
     if (inherits(object$fit, "try-error")) {
       rlang::warn("Model fit failed; cannot make predictions.")

R/predict_time.R

@@ -10,8 +10,7 @@ predict_time.model_fit <- function(object, new_data, ...) {
                             "Use `predict_class()` or `predict_classprob()` for ",
                             "classification models."))
 
-  if (!any(names(object$spec$method$pred) == "time"))
-    rlang::abort("No prediction module defined for this model.")
+  check_spec_pred_type(object, "time")
 
   if (inherits(object$fit, "try-error")) {
     rlang::warn("Model fit failed; cannot make predictions.")