use namespace_case_when() more generally

to close #53

R/model-partykit.R

@@ -12,14 +12,12 @@ orbital.constparty <- function(
 	if (mode == "classification") {
 		res <- character()
 		if ("class" %in% type) {
-      eq <- tidypredict::tidypredict_fit(x)
-      eq <- deparse1(eq)
-      eq <- namespace_case_when(eq)
+			eq <- tidypredict::tidypredict_fit(x)
+			eq <- deparse1(eq)
 			res <- c(res, orbital_tmp_class_name = eq)
 		}
 		if ("prob" %in% type) {
 			eqs <- tidypredict::.extract_partykit_classprob(x)
-			eqs <- namespace_case_when(eqs)
 			names(eqs) <- paste0("orbital_tmp_prob_name", seq_along(lvl))
 			res <- c(res, eqs)
 		}

R/model-xgboost.R

@@ -35,7 +35,6 @@ xgboost_multisoft <- function(x, type, lvl) {
 	trees_split <- split(trees, rep(seq_along(lvl), x$niter))
 	trees_split <- lapply(trees_split, collapse_stumps)
 	trees_split <- vapply(trees_split, paste, character(1), collapse = " + ")
-	trees_split <- namespace_case_when(trees_split)
 
 	res <- stats::setNames(trees_split, lvl)
 
@@ -75,7 +74,6 @@ xgboost_logistic <- function(x, type, lvl) {
 	eq <- tidypredict::tidypredict_fit(x)
 
 	eq <- deparse1(eq)
-	eq <- namespace_case_when(eq)
 
 	res <- NULL
 	if ("class" %in% type) {
@@ -98,12 +96,6 @@ xgboost_logistic <- function(x, type, lvl) {
 	res
 }
 
-namespace_case_when <- function(x) {
-	x <- gsub("dplyr::case_when", "case_when", x)
-	x <- gsub("case_when", "dplyr::case_when", x)
-	x
-}
-
 softmax <- function(lvl) {
 	res <- character(0)
 

R/parsnip.R

@@ -33,6 +33,7 @@ orbital.model_fit <- function(x, ..., prefix = ".pred", type = NULL) {
 		res <- deparse1(res)
 	}
 
+	res <- namespace_case_when(res)
 	res <- set_pred_names(res, x, mode, type, prefix)
 
 	new_orbital_class(res)

R/utils.R

@@ -0,0 +1,7 @@
+namespace_case_when <- function(x) {
+	names <- names(x)
+	x <- gsub("dplyr::case_when", "case_when", x)
+	x <- gsub("case_when", "dplyr::case_when", x)
+	names(x) <- names
+	x
+}

R/workflows.R

@@ -21,5 +21,7 @@ orbital.workflow <- function(x, ..., prefix = ".pred", type = NULL) {
 		attr(out, "pred_names") <- pred_names
 	}
 
+	out <- namespace_case_when(out)
+
 	new_orbital_class(out)
 }