Add vec_slice2() and vec_assign2()

R/slice.R

@@ -249,3 +249,38 @@ vec_slice_seq <- function(x, start, size, increasing = TRUE) {
 vec_slice_rep <- function(x, i, n) {
   .Call(vctrs_slice_rep, x, i, n)
 }
+
+vec_slice2 <- function(x, i) {
+  with_extract(
+    if (vec_is_list(x)) {
+      # Lists are currently guaranteed to have list storage so we can
+      # just subset them directly
+      i <- vec_as_location2(i, vec_size(x))
+      .subset2(x, i)
+    } else {
+      out <- vec_slice(x, i)
+      vec_set_names(out, NULL)
+    }
+  )
+}
+
+vec_assign2 <- function(x, i, value, ..., x_arg = "", value_arg = "") {
+  if (!missing(...)) {
+    ellipsis::check_dots_empty()
+  }
+
+  # We may relax this in the future, e.g. for character `i`
+  if (is_zap(value)) {
+    abort("Can't zap elements.")
+  }
+
+  # If `x` is recursive, wrap RHS in a list before calling
+  # `vec_assign()`. The class of `x` must be coercible with lists. We
+  # intentionally wrap `NULL` values instead of treating them as a
+  # sentinel to zap elements.
+  if (vec_is_list(x)) {
+    value <- list(value)
+  }
+
+  vec_assign(x, i, value, x_arg = x_arg, value_arg = value_arg)
+}

R/subscript.R

@@ -342,3 +342,13 @@ cnd_subscript_scalar <- function(cnd) {
 
   out
 }
+
+with_extract <- function(expr) {
+  withCallingHandlers(
+    vctrs_error_subscript = function(cnd) {
+      cnd$subscript_action <- "extract"
+      cnd_signal(cnd)
+    },
+    expr
+  )
+}

tests/testthat/error/test-slice-assign.txt

@@ -64,3 +64,19 @@ Error: Can't convert `bar` <character> to match type of `foo` <integer>.
 > vec_assign(1:2, 1L, 1:2, value_arg = "bar")
 Error: Can't recycle `bar` (size 2) to size 1.
 
+
+vec_assign2() fails with incompatible type
+==========================================
+
+> vec_assign2(1:3, 2, "")
+Error: Can't convert <character> to <integer>.
+
+
+vec_assign2() fails with OOB subscript
+======================================
+
+> vec_assign2(1:3, 4, 0)
+Error: Can't assign to elements that don't exist.
+x Location 4 doesn't exist.
+i There are only 3 elements.
+

tests/testthat/error/test-slice.txt

@@ -73,3 +73,17 @@ i It must be logical, numeric, or character.
 Error: Must subset elements with a valid subscript vector.
 x Subscript must be a simple vector, not a matrix.
 
+
+vec_slice2() fails if subscript is OOB
+======================================
+
+> vec_slice2(letters, 100)
+Error: Can't extract elements that don't exist.
+x Location 100 doesn't exist.
+i There are only 26 elements.
+
+> vec_slice2(list(), 100)
+Error: Can't extract elements that don't exist.
+x Location 100 doesn't exist.
+i There are only 0 elements.
+

tests/testthat/test-slice-assign.R

@@ -718,6 +718,58 @@ test_that("can assign object of any dimensionality with compact seqs", {
   expect_identical(vec_assign_seq(x4, 2, start, size, increasing), array(rep(c(2, 2, 1), 120), dim = c(3, 4, 5, 6)))
 })
 
+test_that("vec_assign2() handles atomic vectors", {
+  x <- c(a = 1L, b = 2L, c = 3L)
+  exp <- c(a = 1L, b = 0L, c = 3L)
+
+  expect_identical(vec_assign2(x, 2, FALSE), exp)
+
+  local_hidden()
+  expect_identical(vec_assign2(new_hidden(x), 2, FALSE), new_hidden(exp))
+
+  rcrd <- new_rcrd(list(x = 1:3))
+  rcrd_exp <- new_rcrd(list(x = c(1L, 0L, 3L)))
+  expect_identical(vec_assign2(rcrd, 2, new_rcrd(list(x = FALSE))), rcrd_exp)
+})
+
+test_that("vec_assign2() handles lists", {
+  x <- list(a = 1L, b = 2L, c = 3:4)
+  exp1 <- list(a = 1L, b = FALSE, c = 3:4)
+  exp2 <- list(a = 1L, b = NULL, c = 3:4)
+  exp3 <- list(a = 1L, b = list(NULL), c = 3:4)
+
+  expect_identical(vec_assign2(x, 2, FALSE), exp1)
+  expect_identical(vec_assign2(x, 2, NULL), exp2)
+  expect_identical(vec_assign2(x, 2, list(NULL)), exp3)
+
+  local_list_rcrd_methods()
+  expect_identical(vec_assign2(new_list_rcrd(x), 2, FALSE), new_list_rcrd(exp1))
+  expect_identical(vec_assign2(new_list_rcrd(x), 2, NULL), new_list_rcrd(exp2))
+  expect_identical(vec_assign2(new_list_rcrd(x), 2, list(NULL)), new_list_rcrd(exp3))
+})
+
+test_that("zap() is currently disallowed", {
+  expect_error(vec_assign2(list(1), 1, zap()), "Can't zap")
+})
+
+test_that("vec_assign2() fails with incompatible type", {
+  verify_errors({
+    expect_error(
+      vec_assign2(1:3, 2, ""),
+      class = "vctrs_error_incompatible_type"
+    )
+  })
+})
+
+test_that("vec_assign2() fails with OOB subscript", {
+  verify_errors({
+    expect_error(
+      vec_assign2(1:3, 4, 0),
+      class = "vctrs_error_subscript_oob"
+    )
+  })
+})
+
 
 # Golden tests ------------------------------------------------------------
 
@@ -743,5 +795,11 @@ test_that("slice and assign have informative errors", {
     "# `vec_assign()` error args can be overridden"
     vec_assign(1:2, 1L, "x", x_arg = "foo", value_arg = "bar")
     vec_assign(1:2, 1L, 1:2, value_arg = "bar")
+
+    "# vec_assign2() fails with incompatible type"
+    vec_assign2(1:3, 2, "")
+
+    "# vec_assign2() fails with OOB subscript"
+    vec_assign2(1:3, 4, 0)
   })
 })

tests/testthat/test-slice.R

@@ -432,34 +432,6 @@ test_that("vec_slice() works with Altrep classes with custom extract methods", {
   expect_equal(vec_slice(x, idx), c("foo", "foo", "bar"))
 })
 
-test_that("slice has informative error messages", {
-  verify_output(test_path("error", "test-slice.txt"), {
-    "# Unnamed vector with character subscript"
-    vec_slice(1:3, letters[1])
-
-    "# Negative subscripts are checked"
-    vec_slice(1:3, -c(1L, NA))
-    vec_slice(1:3, c(-1L, 1L))
-
-    "# oob error messages are properly constructed"
-    vec_slice(c(bar = 1), "foo")
-
-    "Multiple OOB indices"
-    vec_slice(letters, c(100, 1000))
-    vec_slice(letters, c(1, 100:103, 2, 104:110))
-    vec_slice(set_names(letters), c("foo", "bar"))
-    vec_slice(set_names(letters), toupper(letters))
-
-    "# Can't index beyond the end of a vector"
-    vec_slice(1:2, 3L)
-    vec_slice(1:2, -3L)
-
-    "# vec_slice throws error with non-vector subscripts"
-    vec_slice(1:3, Sys.Date())
-    vec_slice(1:3, matrix(TRUE, ncol = 1))
-  })
-})
-
 # vec_init ----------------------------------------------------------------
 
 test_that("na of atomic vectors is as expected", {
@@ -706,3 +678,101 @@ test_that("column sizes are checked before slicing (#552)", {
   x <- structure(list(a = 1, b = 2:3), row.names = 1:2, class = "data.frame")
   expect_error(vctrs::vec_slice(x, 2), "must match the data frame size")
 })
+
+test_that("vec_slice2() zaps names of atomic values", {
+  expect_identical(
+    vec_slice2(c(foo = 1, bar = 2), 2),
+    2
+  )
+
+  out <- vec_slice2(mtcars, 2)
+  expect_null(vec_names(out))
+  expect_true(vec_equal(out, vec_slice(mtcars, 2)))
+
+  x <- matrix(1:4, 2)
+  row.names(x) <- c("foo", "bar")
+  out <- vec_slice2(x, 2)
+  expect_null(vec_names(out))
+})
+
+test_that("vec_slice2() extracts elements of recursive inputs", {
+  x <- list(a = c(foo = 1), b = c(bar = 2))
+  expect_identical(
+    vec_slice2(x, 2),
+    c(bar = 2)
+  )
+})
+
+test_that("vec_slice2() fails if subscript is OOB", {
+  expect_error(
+    vec_slice2(letters, 100),
+    class = "vctrs_error_subscript_oob"
+  )
+  expect_error(
+    vec_slice2(list(), 100),
+    class = "vctrs_error_subscript_oob"
+  )
+})
+
+test_that("vec_slice2() works with generic atomic vectors", {
+  x <- set_names(new_vctr(1:3), letters[1:3])
+  expect_identical(
+    vec_slice2(x, 2),
+    new_vctr(2L)
+  )
+
+  x <- new_rcrd(list(x = 1:2))
+  expect_identical(
+    vec_slice2(x, 2),
+    new_rcrd(list(x = 2L))
+  )
+})
+
+test_that("vec_slice2() works with generic lists", {
+  x <- list(a = c(foo = 1), b = c(bar = 2))
+  expect_identical(
+    vec_slice2(x, 2),
+    c(bar = 2)
+  )
+
+  local_list_rcrd_methods()
+  expect_identical(
+    vec_slice2(new_list_rcrd(x), 2),
+    c(bar = 2)
+  )
+})
+
+
+# Golden tests -------------------------------------------------------
+
+test_that("slicing functions have informative error messages", {
+  verify_output(test_path("error", "test-slice.txt"), {
+    "# Unnamed vector with character subscript"
+    vec_slice(1:3, letters[1])
+
+    "# Negative subscripts are checked"
+    vec_slice(1:3, -c(1L, NA))
+    vec_slice(1:3, c(-1L, 1L))
+
+    "# oob error messages are properly constructed"
+    vec_slice(c(bar = 1), "foo")
+
+    "Multiple OOB indices"
+    vec_slice(letters, c(100, 1000))
+    vec_slice(letters, c(1, 100:103, 2, 104:110))
+    vec_slice(set_names(letters), c("foo", "bar"))
+    vec_slice(set_names(letters), toupper(letters))
+
+    "# Can't index beyond the end of a vector"
+    vec_slice(1:2, 3L)
+    vec_slice(1:2, -3L)
+
+    "# vec_slice throws error with non-vector subscripts"
+    vec_slice(1:3, Sys.Date())
+    vec_slice(1:3, matrix(TRUE, ncol = 1))
+
+    "# vec_slice2() fails if subscript is OOB"
+    vec_slice2(letters, 100)
+    vec_slice2(list(), 100)
+  })
+})