Correctly support .data pronoun within across (#1685)

Fixes #1520

Author: hadley

NEWS.md

@@ -1,5 +1,6 @@
 # dbplyr (development version)
 
+* `.data$col`, `.data[[col]]`, `.env$var`, and `.env$[[var]]` now work correctly inside `across()` (#1520).
 * New `.sql` pronoun makes it a little easier to use known SQL functions in packages, requiring only `@importFrom dbplyr .sql` (#1117).
 * `join_by(between())` now correctly handles column renames (#1572).
 * SQL Server uses `DATEDIFF_BIG` instead of `DATEDIFF` to work regardless of data size (@edward-burn, #1666).

R/tidyeval-across.R

@@ -210,7 +210,7 @@ across_fun <- function(fun, env, dots, fn) {
       )
     }
 
-    partial_eval_prepare_fun(f_rhs(fun), c(".", ".x"))
+    partial_eval_prepare_fun(f_rhs(fun), c(".", ".x"), env)
   } else if (is_call(fun, "function")) {
     fun <- eval(fun, env)
     partial_eval_fun(fun, env, fn)
@@ -235,10 +235,12 @@ partial_eval_fun <- function(fun, env, fn) {
   }
   args <- fn_fmls_names(fun)
 
-  partial_eval_prepare_fun(body[[2]], args[[1]])
+  partial_eval_prepare_fun(body[[2]], args[[1]], fn_env(fun))
 }
 
-partial_eval_prepare_fun <- function(call, sym) {
+partial_eval_prepare_fun <- function(call, sym, env) {
+  # First resolve any .data/.env pronouns before symbol replacement
+  call <- resolve_mask_pronouns(call, env)
   call <- replace_sym(call, sym, replace = quote(!!.x))
   call <- replace_call(call, replace = quote(!!.cur_col))
   function(x, .cur_col) {
@@ -249,6 +251,31 @@ partial_eval_prepare_fun <- function(call, sym) {
   }
 }
 
+resolve_mask_pronouns <- function(call, env) {
+  if (is_mask_pronoun(call)) {
+    var <- call[[3]]
+
+    if (is_symbol(call[[2]], ".data")) {
+      if (is_call(call, "[[")) {
+        sym(eval(var, env))
+      } else {
+        var
+      }
+    } else {
+      if (is_call(call, "[[")) {
+        env_get(env, var)
+      } else {
+        env_get(env, as.character(var))
+      }
+    }
+  } else if (is_call(call)) {
+    call[] <- lapply(call, resolve_mask_pronouns, env = env)
+    call
+  } else {
+    call
+  }
+}
+
 across_setup <- function(data, call, env, allow_rename, fn, error_call) {
   grps <- group_vars(data)
   tbl <- ungroup(data)

tests/testthat/_snaps/tidyeval-across.md

@@ -80,6 +80,15 @@
       SELECT SUM(`a`) AS `a`, SUM(`b`) AS `b`
       FROM `df`
 
+# lambdas in across() can use columns
+
+    Code
+      show_query(db_across)
+    Output
+      <SQL>
+      SELECT `x` / `y` AS `x`, `y` / `y` AS `y`, `z` / `y` AS `z`
+      FROM `across`
+
 # across() errors if named
 
     Code

tests/testthat/test-tidyeval-across.R

@@ -274,28 +274,48 @@ test_that("across() uses environment from the current quosure (dplyr#5460)", {
 })
 
 test_that("lambdas in across() can use columns", {
-  lf <- lazy_frame(x = 2, y = 4, z = 8)
+  db <- local_memdb_frame("across", x = 2, y = 4, z = 8)
 
   expect_equal(
-    partial_eval_dots(lf, across(everything(), ~ .x / y)),
+    partial_eval_dots(db, across(everything(), ~ .x / y)),
     list(
       x = quo(x / y),
       y = quo(y / y),
       z = quo(z / y)
     )
   )
 
-  skip("not yet correctly supported")
-  # dplyr uses the old value of `y` for division
-  df <- tibble(x = 2, y = 4, z = 8)
-  df |> mutate(across(everything(), ~ .x / .data$y))
-  # so this is the equivalent
-  df |> mutate(data.frame(x = x / y, y = y / y, z = z / y))
-  # dbplyr uses the new value of `y`
-  lf |> mutate(across(everything(), ~ .x / .data$y))
+  db_across <- db |> mutate(across(everything(), ~ .x / y))
+  expect_snapshot(db_across |> show_query())
 
-  # so this is the dbplyr equivalent
-  df |> mutate(x = x / y, y = y / y, z = z / y)
+  # z should be 2 because the value of .data$y is only transformed
+  # _after_ across() is complete, the same as
+  # db |> collect() |> mutate(across(everything(), ~ .x / .data$y))
+  expect_equal(collect(db_across), tibble(x = 0.5, y = 1, z = 2))
+})
+
+test_that("can use .data and .env pronouns(#1520)", {
+  lf <- lazy_frame(x = 1, y = 2)
+
+  my_col <- "y"
+  expect_equal(
+    capture_across(lf, across(x:y, !!quo(~ .x / .data$y))),
+    exprs(x = x / y, y = y / y)
+  )
+  expect_equal(
+    capture_across(lf, across(x:y, !!quo(~ .x / .data[[my_col]]))),
+    exprs(x = x / y, y = y / y)
+  )
+
+  y <- 10
+  expect_equal(
+    capture_across(lf, across(x:y, !!quo(~ .x / .env$y))),
+    exprs(x = x / 10, y = y / 10)
+  )
+  expect_equal(
+    capture_across(lf, across(x:y, !!quo(~ .x / .env[["y"]]))),
+    exprs(x = x / 10, y = y / 10)
+  )
 })
 
 test_that("can pass quosure through `across()`", {