Committed changes that should be updated.

Author: kenmawer

R/step_epi_naomit.R

@@ -1,13 +1,24 @@
-#' Special NA omit step that does two steps in one
+#' Unified NA omission wrapper function for recipes
 #'
-#' @param x Recipe to be used for omission steps
+#' @param recipe Recipe to be used for omission steps
 #'
-#' @return Omits NA's from both predictors and outcomes and training time;
-#' however, only omits predictors at prediction time
+#' @return Omits NA's from both predictors and outcomes at training time;
+#'   however, only omits associated predictors at prediction time to avoid
+#'   losing data.
 #' @export
+#' @examples
+#' tibble(geo_value = rep("place",200),
+#'            time_value = as.Date("2021-01-01") + 0:199,
+#'            case_rate = 1:200,
+#'            death_rate = 1:200) %>%
+#'  as_epi_df() %>%
+#'  recipe() %>%
+#'  step_epi_naomit()
 
-step_epi_naomit <- function(x) {
-  x %>%
+
+step_epi_naomit <- function(recipe) {
+  stopifnot("recipe" %in% class(recipe))
+  recipe %>%
     recipes::step_naomit(all_predictors()) %>%
     recipes::step_naomit(all_outcomes(), skip = TRUE)
 }

man/step_epi_naomit.Rd

@@ -16,6 +16,10 @@ r <- epi_recipe(x) %>%
   step_epi_ahead(death_rate, ahead = 7) %>%
   step_epi_lag(death_rate, lag = c(0,7,14))
 
+test_that("Argument must be a recipe", {
+  expect_error(step_epi_naomit(x))
+})
+
 z1 <- step_epi_naomit(r)
 z2 <- r %>%
   step_naomit(all_predictors()) %>%