Keep the epi_df class when head and tail are used

DESCRIPTION

@@ -37,7 +37,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/epi_df.R

@@ -220,9 +220,9 @@ as_epi_df.tbl_df = function(x, geo_type, time_type, as_of,
   if (!("time_value" %in% names(x))) {
     Abort("`x` must contain a `time_value` column.")
   }
-
+  
   new_epi_df(x, geo_type, time_type, as_of,
-             additional_metadata = list(), ...)
+             additional_metadata, ...)
 }
 
 #' @method as_epi_df data.frame
@@ -259,4 +259,4 @@ as_epi_df.tbl_ts = function(x, geo_type, time_type, as_of,
 #' @export
 is_epi_df = function(x) {
   inherits(x, "epi_df")
-}
+}

R/methods-epi_df.R

@@ -62,20 +62,62 @@ 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), ...)
-}
+`[.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))
 
-#' @method tail epi_df
-#' @importFrom utils tail
-#' @export
-#' @noRd
-tail.epi_df = function(x, ...) {
-  tail(tibble::as_tibble(x), ...)
+  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 = attr(x, "metadata")
+  new_epi_df(tibble::as_tibble(res), 
+             geo_type = att_x$geo_type, 
+             time_type = att_x$time_type, 
+             as_of = att_x$as_of,
+             additional_metadata = 
+               att_x[!(names(att_x)  %in% c("geo_type", "time_type", "as_of"))])
 }
 
 #' `dplyr` verbs

tests/testthat/test-methods-epi_df.R

@@ -0,0 +1,73 @@
+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),
+  indicator_var = as.factor(rep(1:2, times = 5)), 
+) %>% as_epi_df(additional_metadata = c(other_keys = "indicator_var"))
+
+att_toy = attr(toy_epi_df, "metadata")
+
+test_that("head and tail do not drop the epi_df class", {
+  att_head = attr(head(toy_epi_df), "metadata")
+  att_tail = attr(tail(toy_epi_df), "metadata")
+
+  expect_true(is_epi_df(head(toy_epi_df)))
+  expect_true(is_epi_df(tail(toy_epi_df)))
+  expect_identical(att_head$geo_type, att_toy$geo_type)
+  expect_identical(att_head$time_type, att_toy$time_type)
+  expect_identical(att_head$as_of, att_toy$as_of)
+  expect_identical(att_head$other_keys, att_toy$other_keys)
+  expect_identical(att_tail$geo_type, att_toy$geo_type)
+  expect_identical(att_tail$time_type, att_toy$time_type)
+  expect_identical(att_tail$as_of, att_toy$as_of)
+  expect_identical(att_tail$other_keys, att_toy$other_keys)
+})
+
+
+test_that("subsetting drops or does not drop the epi_df class appropriately", {
+
+  # Row subset - should be epi_df
+  row_subset = toy_epi_df[1:2, ]
+  att_row_subset = attr(row_subset, "metadata")
+
+  expect_true(is_epi_df(row_subset))
+  expect_equal(nrow(row_subset), 2L)
+  expect_equal(ncol(row_subset), 5L)
+  expect_identical(att_row_subset$geo_type, att_toy$geo_type)
+  expect_identical(att_row_subset$time_type, att_toy$time_type)
+  expect_identical(att_row_subset$as_of, att_toy$as_of)
+  expect_identical(att_row_subset$other_keys, att_toy$other_keys)
+
+  # Col subset - shouldn't be an epi_df
+  col_subset = toy_epi_df[, 2:3]
+
+  expect_false(is_epi_df(col_subset))
+  expect_true(tibble::is_tibble(col_subset))
+  expect_equal(nrow(col_subset), 10L)
+  expect_equal(ncol(col_subset), 2L)
+
+  # Row and col single value - shouldn't be an epi_df
+  row_col_subset1 = toy_epi_df[1,2]
+  expect_false(is_epi_df(row_col_subset1))
+  expect_true(tibble::is_tibble(row_col_subset1))
+  expect_equal(nrow(row_col_subset1), 1L)
+  expect_equal(ncol(row_col_subset1), 1L)
+
+  # Row and col subset that contains geo_value and time_value - should be epi_df
+  row_col_subset2 = toy_epi_df[2:3,1:3]
+  att_row_col_subset2 = attr(row_col_subset2, "metadata")
+
+  expect_true(is_epi_df(row_col_subset2))
+  expect_equal(nrow(row_col_subset2), 2L)
+  expect_equal(ncol(row_col_subset2), 3L)
+  expect_identical(att_row_col_subset2$geo_type, att_toy$geo_type)
+  expect_identical(att_row_col_subset2$time_type, att_toy$time_type)
+  expect_identical(att_row_col_subset2$as_of, att_toy$as_of)
+  expect_identical(att_row_col_subset2$other_keys, att_toy$other_keys)
+
+})