Factor out eval_relocate() and use dplyr_col_select()

R/relocate.R

@@ -52,49 +52,103 @@ relocate <- function(.data, ..., .before = NULL, .after = NULL) {
 
 #' @export
 relocate.data.frame <- function(.data, ..., .before = NULL, .after = NULL) {
-  to_move <- tidyselect::eval_select(expr(c(...)), .data)
+  loc <- eval_relocate(
+    expr = expr(c(...)),
+    data = .data,
+    before = enquo(.before),
+    after = enquo(.after),
+    before_arg = ".before",
+    after_arg = ".after"
+  )
 
-  # Enforce the invariant that `ncol(.data) == ncol(relocate(.data, ...))` by
+  dplyr_col_select(.data, loc, names = names(loc))
+}
+
+eval_relocate <- function(expr,
+                          data,
+                          ...,
+                          before = NULL,
+                          after = NULL,
+                          before_arg = "before",
+                          after_arg = "after",
+                          env = caller_env(),
+                          error_call = caller_env()) {
+  # `eval_relocate()` returns a named integer vector of size `ncol(data)`
+  # describing how to rearrange `data`. Each location in the range
+  # `seq2(1L, ncol(data))` is represented once. The names are the new names to
+  # assign to those columns. They are typically the same as the original names,
+  # but `expr` does allow for renaming.
+
+  check_dots_empty0(...)
+
+  sel <- tidyselect::eval_select(
+    expr = expr,
+    data = data,
+    env = env,
+    error_call = error_call
+  )
+
+  # Enforce the invariant that relocating can't change the number of columns by
   # retaining only the last instance of a column that is renamed multiple times
-  # while it is being moved
   # TODO: https://github.com/r-lib/vctrs/issues/1442
-  # `to_move <- vec_unique(to_move, which = "last")`
-  loc_last <- which(!duplicated(to_move, fromLast = TRUE))
-  to_move <- vec_slice(to_move, loc_last)
+  # `sel <- vec_unique(sel, which = "last")`
+  loc_last <- which(!duplicated(sel, fromLast = TRUE))
+  sel <- vec_slice(sel, loc_last)
 
-  .before <- enquo(.before)
-  .after <- enquo(.after)
-  has_before <- !quo_is_null(.before)
-  has_after <- !quo_is_null(.after)
+  n <- length(data)
+
+  before <- as_quosure(before, env = env)
+  after <- as_quosure(after, env = env)
+
+  has_before <- !quo_is_null(before)
+  has_after <- !quo_is_null(after)
 
   if (has_before && has_after) {
-    abort("Must supply only one of `.before` and `.after`.")
-  } else if (has_before) {
-    where <- min(unname(tidyselect::eval_select(.before, .data)))
-    if (!where %in% to_move) {
-      to_move <- c(to_move, where)
+    message <- glue("Can't supply both `{before_arg}` and `{after_arg}`.")
+    abort(message, call = error_call)
+  }
+
+  if (has_before) {
+    # TODO: Use `allow_rename = FALSE`. https://github.com/r-lib/tidyselect/issues/225
+    where <- tidyselect::eval_select(before, data, env = env, error_call = error_call)
+    where <- unname(where)
+
+    if (length(where) == 0L) {
+      # Empty `before` selection pushes `sel` to the front
+      where <- 1L
+    } else {
+      where <- min(where)
     }
   } else if (has_after) {
-    where <- max(unname(tidyselect::eval_select(.after, .data)))
-    if (!where %in% to_move) {
-      to_move <- c(where, to_move)
+    # TODO: Use `allow_rename = FALSE`. https://github.com/r-lib/tidyselect/issues/225
+    where <- tidyselect::eval_select(after, data, env = env, error_call = error_call)
+    where <- unname(where)
+
+    if (length(where) == 0L) {
+      # Empty `after` selection pushes `sel` to the back
+      where <- n
+    } else {
+      where <- max(where)
     }
+
+    where <- where + 1L
   } else {
+    # Defaults to `before = everything()` if neither `before` nor `after` are supplied
     where <- 1L
-    if (!where %in% to_move) {
-      to_move <- c(to_move, where)
-    }
   }
 
-  lhs <- setdiff(seq2(1, where - 1), to_move)
-  rhs <- setdiff(seq2(where + 1, ncol(.data)), to_move)
+  lhs <- seq2(1L, where - 1L)
+  rhs <- seq2(where, n)
 
-  pos <- c(lhs, to_move, rhs)
-  out <- .data[pos]
-  new_names <- names(pos)
+  lhs <- setdiff(lhs, sel)
+  rhs <- setdiff(rhs, sel)
 
-  if (!is.null(new_names)) {
-    names(out)[new_names != ""] <- new_names[new_names != ""]
-  }
-  out
+  names <- names(data)
+
+  names(lhs) <- names[lhs]
+  names(rhs) <- names[rhs]
+
+  sel <- vec_c(lhs, sel, rhs)
+
+  sel
 }

tests/testthat/_snaps/relocate.md

@@ -0,0 +1,8 @@
+# can only supply one of .before and .after
+
+    Code
+      relocate(df, .before = 1, .after = 1)
+    Condition
+      Error in `relocate()`:
+      ! Can't supply both `.before` and `.after`.
+

tests/testthat/test-mutate.r

@@ -372,6 +372,17 @@ test_that("can use .before and .after to control column position", {
   expect_named(mutate(df, x = 1, .after = y), c("x", "y"))
 })
 
+test_that("attributes of bare data frames are retained when `.before` and `.after` are used (#6341)", {
+  # We require `[` methods to be in charge of keeping extra attributes for all
+  # data frame subclasses (except for data.tables)
+  df <- vctrs::data_frame(x = 1, y = 2)
+  attr(df, "foo") <- "bar"
+
+  out <- mutate(df, z = 3, .before = x)
+
+  expect_identical(attr(out, "foo"), "bar")
+})
+
 test_that(".keep and .before/.after interact correctly", {
   df <- tibble(x = 1, y = 1, z = 1, a = 1, b = 2, c = 3) %>%
     group_by(a, b)

tests/testthat/test-relocate.R

@@ -1,3 +1,6 @@
+# ------------------------------------------------------------------------------
+# relocate()
+
 test_that(".before and .after relocate individual cols", {
   df <- tibble(x = 1, y = 2)
   expect_named(relocate(df, x, .after = y), c("y", "x"))
@@ -23,7 +26,10 @@ test_that("no .before/.after moves to front", {
 
 test_that("can only supply one of .before and .after", {
   df <- tibble(x = 1)
-  expect_error(relocate(df, .before = 1, .after = 1), "only one")
+
+  expect_snapshot(error = TRUE, {
+    relocate(df, .before = 1, .after = 1)
+  })
 })
 
 test_that("before and after are defused with context", {
@@ -91,3 +97,109 @@ test_that("`relocate()` retains the last duplicate when renaming while moving (#
     select(rename(df, a = x, b = y, c = x), b, c)
   )
 })
+
+test_that("attributes of bare data frames are retained (#6341)", {
+  # We require `[` methods to be in charge of keeping extra attributes for all
+  # data frame subclasses (except for data.tables)
+  df <- vctrs::data_frame(x = 1, y = 2)
+  attr(df, "foo") <- "bar"
+
+  out <- relocate(df, y, .before = x)
+
+  expect_identical(attr(out, "foo"), "bar")
+})
+
+# ------------------------------------------------------------------------------
+# eval_relocate()
+
+test_that("works with zero column data frames (#6167)", {
+  data <- tibble()
+  expr <- expr(any_of("b"))
+
+  expect_identical(
+    eval_relocate(expr, data),
+    set_names(integer())
+  )
+})
+
+test_that("works with `before` and `after` `everything()`", {
+  data <- tibble(w = 1, x = 2, y = 3, z = 4)
+  expr <- expr(c(y, z))
+  expr_everything <- expr(everything())
+
+  expect_identical(
+    eval_relocate(expr, data, before = expr_everything),
+    c(y = 3L, z = 4L, w = 1L, x = 2L)
+  )
+  expect_identical(
+    eval_relocate(expr, data, after = expr_everything),
+    c(w = 1L, x = 2L, y = 3L, z = 4L)
+  )
+})
+
+test_that("moves columns to the front when neither `before` nor `after` are specified", {
+  data <- tibble(x = 1, y = 2, z = 3)
+  expr <- expr(c(z, y))
+
+  expect_identical(
+    eval_relocate(expr, data),
+    c(z = 3L, y = 2L, x = 1L)
+  )
+})
+
+test_that("Empty `before` selection moves columns to front", {
+  data <- tibble(x = 1, y = 2, z = 3)
+  expr <- expr(y)
+  before <- expr(where(is.character))
+
+  expect_identical(
+    eval_relocate(expr, data, before = before),
+    c(y = 2L, x = 1L, z = 3L)
+  )
+})
+
+test_that("Empty `after` selection moves columns to end", {
+  data <- tibble(x = 1, y = 2, z = 3)
+  expr <- expr(y)
+  after <- expr(where(is.character))
+
+  expect_identical(
+    eval_relocate(expr, data, after = after),
+    c(x = 1L, z = 3L, y = 2L)
+  )
+})
+
+test_that("Empty `before` and `after` selections work with 0-col data frames", {
+  data <- tibble()
+  expr <- expr(any_of("a"))
+  expr_is_character <- expr(where(is.character))
+
+  expect_identical(
+    eval_relocate(expr, data, before = expr_is_character),
+    set_names(integer())
+  )
+
+  expect_identical(
+    eval_relocate(expr, data, after = expr_is_character),
+    set_names(integer())
+  )
+})
+
+test_that("retains the last duplicate when renaming while moving (#6209)", {
+  # To enforce the invariant that relocating can't change the number of columns
+  data <- tibble(x = 1)
+  expr <- expr(c(a = x, b = x))
+
+  expect_identical(
+    eval_relocate(expr, data),
+    c(b = 1L)
+  )
+
+  data <- tibble(x = 1, y = 2)
+  expr <- expr(c(a = x, b = y, c = x))
+
+  expect_identical(
+    eval_relocate(expr, data),
+    c(b = 2L, c = 1L)
+  )
+})