Handle lists correctly in nth.()

NEWS.md

@@ -14,6 +14,7 @@
 * `across.()`: Can namespace functions in `.fns` arg (#511)
 * `as_tidytable()`: Can keep row names when converting a matrix (#527)
 * `unnest.()`: Handles empty data frames (@roboton, #530)
+* `nth.()`: Extracts list elements (#534)
 
 # tidytable 0.8.0
 

R/first-last-nth.R

@@ -18,19 +18,19 @@
 #' first.(vec)
 #' last.(vec)
 #' nth.(vec, 4)
-first. <- function(x, default = NA, na_rm = FALSE) {
+first. <- function(x, default = NULL, na_rm = FALSE) {
   nth.(x, 1L, default, na_rm)
 }
 
 #' @rdname first.
 #' @export
-last. <- function(x, default = NA, na_rm = FALSE) {
+last. <- function(x, default = NULL, na_rm = FALSE) {
   nth.(x, -1L, default, na_rm)
 }
 
 #' @rdname first.
 #' @export
-nth. <- function(x, n, default = NA, na_rm = FALSE) {
+nth. <- function(x, n, default = NULL, na_rm = FALSE) {
   if (na_rm) {
     x <- x[!vec_equal_na(x)]
   }
@@ -41,7 +41,27 @@ nth. <- function(x, n, default = NA, na_rm = FALSE) {
   }
 
   if (n > size || n == 0) {
-    vec_cast(default, vec_ptype(x))
+    nth_default(x, default)
+  } else {
+    vec_slice2(x, n)
+  }
+}
+
+nth_default <- function(x, default) {
+  if (vec_is_list(x)) {
+    out <- default
+  } else {
+    if (is.null(default)) {
+      default <- NA
+    }
+    out <- vec_cast(default, vec_ptype(x))
+  }
+  out
+}
+
+vec_slice2 <- function(x, n) {
+  if (vec_is_list(x)) {
+    .subset2(x, n)
   } else {
     vec_slice(x, n)
   }

tests/testthat/test-first-last-nth.R

@@ -17,3 +17,10 @@ test_that("work on data frames", {
   expect_equal(first.(df), head(df, 1))
   expect_equal(last.(df), tail(df, 1))
 })
+
+test_that("work on lists", {
+  l <- list(x = "x", y = "y")
+  expect_equal(first.(l), "x")
+  expect_equal(last.(l), "y")
+  expect_equal(nth.(l, 3), NULL)
+})