Keep the epi_df class when head and tail are used

DESCRIPTION

@@ -36,7 +36,8 @@ Imports:
     tidyr,
     tidyselect,
     tsibble,
-    utils
+    utils,
+    vctrs
 Suggests:
     covidcast,
     delphi.epidata,

NAMESPACE

@@ -9,14 +9,12 @@ S3method(as_tsibble,epi_df)
 S3method(filter,epi_df)
 S3method(group_by,epi_df)
 S3method(group_modify,epi_df)
-S3method(head,epi_df)
 S3method(mutate,epi_df)
 S3method(print,epi_df)
 S3method(relocate,epi_df)
 S3method(rename,epi_df)
 S3method(slice,epi_df)
 S3method(summary,epi_df)
-S3method(tail,epi_df)
 S3method(ungroup,epi_df)
 S3method(unnest,epi_df)
 export("%>%")

R/methods-epi_df.R

@@ -62,20 +62,56 @@ summary.epi_df = function(object, ...) {
                          dplyr::summarize(mean(.data$num)))))
 }
 
-#' @method head epi_df
-#' @importFrom utils head
 #' @export
-#' @noRd
-head.epi_df = function(x, ...) {
-  head(tibble::as_tibble(x), ...)
-}
-
-#' @method tail epi_df
-#' @importFrom utils tail
-#' @export
-#' @noRd
-tail.epi_df = function(x, ...) {
-  tail(tibble::as_tibble(x), ...)
+`[.epi_df` <- function(x, i, j, drop = FALSE) { 
+  res <- NextMethod() 
+
+  if (!is.data.frame(res)) return(res)
+
+  i_arg <- substitute(i) 
+  j_arg <- substitute(j) 
+
+  if (missing(i)) {
+    i <- NULL
+    i_arg <- NULL
+  } else if (is.null(i)) {
+    i <- integer()
+  }
+
+  if (missing(j)) {
+    j <- NULL
+    j_arg <- NULL
+  } else if (is.null(j)) {
+    j <- integer()
+  }
+
+  # Ignore drop as an argument for counting
+  n_real_args <- nargs() - !missing(drop) 
+
+  # Case when the number of args (excluding drop) is not 3 or more
+  if (n_real_args <= 2L) {
+    j <- i
+    i <- NULL
+    j_arg <- i_arg
+    i_arg <- NULL
+  }
+
+  cn <- names(res)
+  nr <- vctrs::vec_size(x) 
+  not_epi_df <- !("time_value" %in% cn) || !("geo_value" %in% cn) || vctrs::vec_size(res) > nr || any(i > nr)
+  if (not_epi_df) return(tibble::as_tibble(res))
+
+  # Case when i is numeric and there are duplicate values in it
+  if (is.numeric(i) && vctrs::vec_duplicate_any(i) > 0) return(tibble::as_tibble(res))
+
+  # Column subsetting only, then return res as tibble
+  if (rlang::is_null(i) && !rlang::is_null(j)) return(tibble::as_tibble(res))
+
+  att_x = attributes(x)
+  new_epi_df(tibble::as_tibble(res), geo_type = att_x$metadata$geo_type, 
+             time_type = att_x$metadata$time_type, as_of = att_x$metadata$as_of,
+             additional_metadata = 
+               as.list(names(att_x$metadata)[!(names(att_x$metadata) %in% c("geo_type", "time_type", "as_of"))]))
 }
 
 #' `dplyr` verbs

tests/testthat/test-methods-epi_df.R

@@ -0,0 +1,21 @@
+test_that("head and tail do not drop the epi_df class", {
+  toy_epi_df <- tibble::tibble(
+    x = 1:10,
+    y = 1:10,
+    time_value = rep(seq(
+      as.Date("2020-01-01"),
+      by = 1,
+      length.out = 5
+    ), times = 2),
+    geo_value = rep(c("ca", "hi"), each = 5)
+  ) %>% epiprocess::as_epi_df()
+
+  expect_true(is_epi_df(head(toy_epi_df)))
+  expect_true(is_epi_df(tail(toy_epi_df)))
+  expect_identical(attributes(head(toy_epi_df))$metadata$geo_type, attributes(toy_epi_df)$metadata$geo_type)
+  expect_identical(attributes(head(toy_epi_df))$metadata$time_type, attributes(toy_epi_df)$metadata$time_type)
+  expect_identical(attributes(head(toy_epi_df))$metadata$as_of, attributes(toy_epi_df)$metadata$as_of)
+  expect_identical(attributes(tail(toy_epi_df))$metadata$geo_type, attributes(toy_epi_df)$metadata$geo_type)
+  expect_identical(attributes(tail(toy_epi_df))$metadata$time_type, attributes(toy_epi_df)$metadata$time_type)
+  expect_identical(attributes(tail(toy_epi_df))$metadata$as_of, attributes(toy_epi_df)$metadata$as_of)
+})