one_hot_encoder() can return the new variables as integer using the as_integer parameter

Author: aldomann

R/one_hot_encoder.R

@@ -1,15 +1,16 @@
 #' One-hot Encoder
 #'
 #' @param data Input data frame
-#' @param feature Unquoted of the feature/column to encode
+#' @param feature Unquoted form of the feature/column to encode
+#' @param as_integer Wether the boolean should be converted to integer or not
 #'
 #' @return
 #' @export
 #'
 #' @examples
 #' set.seed(11)
 #' one_hot_encoder(iris[sample(1:150, 10),], Species)
-one_hot_encoder <- function(data, feature) {
+one_hot_encoder <- function(data, feature, as_integer = FALSE) {
 
   # Deparse feature name
   feature_name <- deparse(substitute(feature))
@@ -44,5 +45,13 @@ one_hot_encoder <- function(data, feature) {
     data[[labels[i]]] <- stringr::str_detect(processed_feature, detection_labels[i])
   }
 
+  if (as_integer) {
+    data <- data %>%
+      dplyr::mutate_at(
+        .vars = dplyr::vars(labels),
+        .funs = as.integer
+      )
+  }
+
   return(data %>% dplyr::select(-{{ feature }}))
 }