Translate pick()

tidyverse · Nov 10, 2022 · f1bd15b · f1bd15b
1 parent 5d71f4c
commit f1bd15b
Show file tree

Hide file tree

Showing 2 changed files with 28 additions and 13 deletions.
diff --git a/R/tidyeval-across.R b/R/tidyeval-across.R
@@ -3,9 +3,13 @@ capture_across <- function(data, x, j = TRUE) {
   dt_squash_across(get_expr(x), get_env(x), data, j)
 }
 
-dt_squash_across <- function(call, env, data, j = j) {
+dt_squash_across <- function(call, env, data, j = j, is_top_across = TRUE) {
   call <- match.call(dplyr::across, call, expand.dots = FALSE, envir = env)
-  across_setup(data, call, env, allow_rename = TRUE, j = j, fn = "across()")
+  out <- across_setup(data, call, env, allow_rename = TRUE, j = j, fn = "across()")
+  if (is_false(is_top_across)) {
+    out <- call2("data.table", !!!out)
+  }
+  out
 }
 
 capture_if_all <- function(data, x, j = TRUE) {

diff --git a/R/tidyeval.R b/R/tidyeval.R
@@ -30,7 +30,8 @@ globalVariables(dt_funs)
 
 capture_dots <- function(.data, ..., .j = TRUE) {
   dots <- enquos(..., .named = .j)
-  dots <- lapply(dots, dt_squash, data = .data, j = .j)
+  top_across <- map(dots, quo_is_call, "across")
+  dots <- map2(dots, top_across, ~ dt_squash(.x, data = .data, j = .j, is_top_across = .y))
 
   # Remove names from any list elements
   is_list <- map_lgl(dots, is.list)
@@ -44,7 +45,8 @@ capture_dots <- function(.data, ..., .j = TRUE) {
 capture_new_vars <- function(.data, ...) {
   dots <- as.list(enquos(..., .named = TRUE))
   for (i in seq_along(dots)) {
-    dot <- dt_squash(dots[[i]], data = .data)
+    dot <- dots[[i]]
+    dot <- dt_squash(dot, data = .data, is_top_across = quo_is_call(dot, "across"))
     if (is.null(dot)) {
       dots[i] <- list(NULL)
     } else {
@@ -67,7 +69,7 @@ capture_dot <- function(.data, x, j = TRUE) {
 }
 
 # squash quosures
-dt_squash <- function(x, env, data, j = TRUE) {
+dt_squash <- function(x, env, data, j = TRUE, is_top_across = TRUE) {
   if (is_atomic(x) || is_null(x)) {
     x
   } else if (is_symbol(x)) {
@@ -101,21 +103,30 @@ dt_squash <- function(x, env, data, j = TRUE) {
       }
     }
   } else if (is_quosure(x)) {
-    dt_squash(get_expr(x), get_env(x), data, j = j)
+    dt_squash(get_expr(x), get_env(x), data, j = j, is_top_across)
   } else if (is_call(x, "if_any")) {
     dt_squash_if(x, env, data, j = j, reduce = "|")
   } else if (is_call(x, "if_all")) {
     dt_squash_if(x, env, data, j = j, reduce = "&")
   } else if (is_call(x, "across")) {
-    dt_squash_across(x, env, data, j = j)
+    dt_squash_across(x, env, data, j = j, is_top_across)
+  } else if (is_call(x, "pick")) {
+    call <- call_match(x, pick, dots_expand = FALSE)
+    .cols <- call2("c", !!!call$...)
+    across_call <- call2("across", .cols)
+    dt_squash_across(across_call, env, data, j, is_top_across)
   } else if (is_call(x)) {
-    dt_squash_call(x, env, data, j = j)
+    dt_squash_call(x, env, data, j = j, is_top_across)
   } else {
     abort("Invalid input")
   }
 }
 
-dt_squash_call <- function(x, env, data, j = TRUE) {
+pick <- function(...) {
+  "yep"
+}
+
+dt_squash_call <- function(x, env, data, j = TRUE, is_top_across = TRUE) {
   if (is_mask_pronoun(x)) {
     var <- x[[3]]
     if (is_call(x, "[[")) {
@@ -127,7 +138,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
       sym(paste0("..", var))
     }
   } else if (is_call(x, c("coalesce", "replace_na"))) {
-    args <- lapply(x[-1], dt_squash, env = env, data = data, j = j)
+    args <- lapply(x[-1], dt_squash, env, data, j, is_top_across)
     call2("fcoalesce", !!!args)
   } else if (is_call(x, "case_when")) {
     # case_when(x ~ y) -> fcase(x, y)
@@ -139,7 +150,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
         x[[3]]
       )
     }))
-    args <- lapply(args, dt_squash, env = env, data = data, j = j)
+    args <- lapply(args, dt_squash, env = env, data = data, j = j, is_top_across)
     call2("fcase", !!!args)
   } else if (is_call(x, "cur_data")) {
     quote(.SD)
@@ -164,7 +175,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
     }
 
     x[[1]] <- quote(fifelse)
-    x[-1] <- lapply(x[-1], dt_squash, env, data, j = j)
+    x[-1] <- lapply(x[-1], dt_squash, env, data, j = j, is_top_across)
     x
   } else if (is_call(x, c("lag", "lead"))) {
     if (is_call(x, "lag")) {
@@ -237,7 +248,7 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
     }
     call
   } else {
-    x[-1] <- lapply(x[-1], dt_squash, env, data, j = j)
+    x[-1] <- lapply(x[-1], dt_squash, env, data, j = j, is_top_across)
     x
   }
 }