add keep argument to dt_unnest

Closes #38

R/unnest.R

@@ -3,7 +3,8 @@
 #' Quickly unnest data tables, particularly those nested by \code{dt_nest()}.
 #'
 #' @param dt_ the data table to unnest
-#' @param col  the column to unnest
+#' @param col the column to unnest
+#' @param keep whether to keep the nested column, default is \code{TRUE}
 #'
 #' @examples
 #'
@@ -19,14 +20,14 @@
 #' @import data.table
 #'
 #' @export
-dt_unnest <- function(dt_, col) {
+dt_unnest <- function(dt_, col, keep = TRUE) {
   UseMethod("dt_unnest", dt_)
 }
 
 #' @export
-dt_unnest.default <- function(dt_, col) {
-  if (isFALSE(is.data.table(dt_))) {
-    dt_ <- as.data.table(dt_)
+dt_unnest.default <- function(dt_, col, keep = TRUE) {
+  if (isFALSE(data.table::is.data.table(dt_))) {
+    dt_ <- data.table::as.data.table(dt_)
   }
 
   # col to unnest
@@ -45,7 +46,13 @@ dt_unnest.default <- function(dt_, col) {
   others <- others[!others_class %in% c("list", "data.table", "data.frame", "tbl_df")]
 
   # Join them all together
-  dt_[seq_len(.N), eval(col)[[1L]], by = others][dt_, on = others]
+  dt_2 <- dt_[seq_len(.N), eval(col)[[1L]], by = others]
+
+  if (isFALSE(keep)){
+    dt_[, deparse(col) := NULL]
+  }
+
+  dt_2[dt_, on = others]
 }
 
 

tests/testthat/test-dt_nest.R

@@ -5,7 +5,7 @@ starwars <- dplyr::starwars
 
 test_that("dt_nest() works", {
   res <- dt_nest(starwars, species, homeworld)
-  expect_is(dplyr::pull(res), "list")
+  expect_type(dplyr::pull(res), "list")
 
   dplyr_res <- dplyr::group_nest(starwars, species, homeworld)
 
@@ -107,6 +107,23 @@ test_that("can unnest mixture of name and unnamed lists of same length", {
   )
 })
 
+
+test_that("unnest keep argument", {
+  test_df <- data.table::data.table(x = 1:2)
+  nested_df <- data.table::data.table(id = 1L:2L, list_column = list(test_df, test_df))
+  expect_identical(
+    dt_unnest(nested_df, list_column, keep = TRUE),
+    data.table::data.table(id = c(1L,1L,2L,2L), x = c(1L,2L,1L,2L), list_column = list(test_df, test_df))
+  )
+
+  expect_identical(
+    dt_unnest(nested_df, list_column, keep = FALSE),
+    data.table::data.table(id = c(1L,1L,2L,2L), x = c(1L,2L,1L,2L))
+  )
+})
+
+
+
 # needs to be implemented in dt_hoist()
 # test_that("hoist: elements must all be of same type", {
 #   df <- data.table::data.table(x = list(1, "a"),