respect predict(type) with a postprocessor

Closes #251, closes #234. Another route we could have taken here is to take a `type` in `predict.tailor()`, but this would lead to a different type needing to be supplied to `predict.model_fit()` than `predict.tailor()` in interactive usage if the user is only interested in hard class predictions.

R/predict.R

@@ -73,17 +73,33 @@ predict.workflow <- function(object, new_data, type = NULL, opts = list(), ...)
     return(predict(fit, new_data, type = type, opts = opts, ...))
   }
 
-  # use `augment()` rather than `fit()` to get all possible prediction `type`s.
-  # likely, we actually want tailor to check for the existence of needed
-  # columns at predict time and just use `predict()` output here.
+  # use `augment()` rather than `fit()` to get all possible prediction `type`s (#234).
   fit_aug <- augment(fit, new_data, opts = opts, ...)
 
   post <- extract_postprocessor(workflow)
-  predict(post, fit_aug)[post$columns$estimate]
+  predict(post, fit_aug)[predict_type_column_names(type, post$columns)]
 }
 
 forge_predictors <- function(new_data, workflow) {
   mold <- extract_mold(workflow)
   forged <- hardhat::forge(new_data, blueprint = mold$blueprint)
   forged$predictors
 }
+
+predict_type_column_names <- function(type, tailor_columns, call = caller_env()) {
+  check_string(type, allow_null = TRUE, call = call)
+
+  if (is.null(type)) {
+    return(tailor_columns$estimate)
+  }
+
+  switch(
+    type,
+    numeric = , class = tailor_columns$estimate,
+    prob = tailor_columns$probabilities,
+    cli::cli_abort(
+      "Unsupported prediction {.arg type} {.val {type}} for a workflow with a postprocessor.",
+      call = call
+    )
+  )
+}

tests/testthat/_snaps/predict.md

@@ -7,3 +7,11 @@
       ! Can't predict on an untrained workflow.
       i Do you need to call `fit()`?
 
+# predict(type) is respected with a postprocessor (#251)
+
+    Code
+      predict(wflow_fit, d[1:5, ], type = "boop")
+    Condition
+      Error in `predict()`:
+      ! Unsupported prediction `type` "boop" for a workflow with a postprocessor.
+

tests/testthat/test-predict.R

@@ -155,3 +155,21 @@ test_that("monitoring: no double intercept due to dot expansion in model formula
   # so lm()'s predict method won't error anymore here. (tidymodels/parsnip#1033)
   expect_no_error(predict(fit_with_intercept, mtcars))
 })
+
+test_that("predict(type) is respected with a postprocessor (#251)", {
+  # create example data
+  y <- seq(0, 7, .1)
+  d <- data.frame(y = as.factor(ifelse(y > 3.5, "yes", "no")), x = y + (y-3)^2)
+  wflow <- workflow(y ~ ., parsnip::logistic_reg(), tailor::tailor())
+  wflow_fit <- fit(wflow, d)
+
+  pred_class <- predict(wflow_fit, d[1:5,], type = "class")
+  pred_prob <- predict(wflow_fit, d[1:5,], type = "prob")
+  pred_null <- predict(wflow_fit, d[1:5,])
+
+  expect_named(pred_class, ".pred_class")
+  expect_named(pred_prob, c(".pred_no", ".pred_yes"), ignore.order = TRUE)
+  expect_equal(pred_class, pred_null)
+
+  expect_snapshot(error = TRUE, predict(wflow_fit, d[1:5,], type = "boop"))
+})