From 9c83443be6cdef84bb7b0d97012f159cc6e3bcb4 Mon Sep 17 00:00:00 2001
From: Mark Fairbanks <mark.t.fairbanks@gmail.com>
Date: Mon, 8 Mar 2021 16:54:37 -0700
Subject: [PATCH] Use dt_squash on if_else args (#221)

* Use dt_squash on if_else args. Translate ifelse to fifelse. Fixes #220.

* Update tests

* Update relocate docs
---
 NEWS.md                        |  4 ++++
 R/tidyeval.R                   |  5 +++--
 man/relocate.dtplyr_step.Rd    |  2 +-
 tests/testthat/test-tidyeval.R | 24 +++++++++++++++++++-----
 4 files changed, 27 insertions(+), 8 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 262092826..3e66b9f14 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -18,6 +18,10 @@
 
 * `tally()` is now translated (@mgirlich, #201).
 
+* `ifelse()` is now mapped to `fifelse()` (@markfairbanks, #220).
+
+* `.data` and `.env` pronouns now work inside of `if_else()` (@markfairbanks, #220).
+
 * More tidyr verbs are in the process of being added:
 
   * `drop_na()` (@markfairbanks, #194)
diff --git a/R/tidyeval.R b/R/tidyeval.R
index d4cdf6d6f..b97f29545 100644
--- a/R/tidyeval.R
+++ b/R/tidyeval.R
@@ -15,7 +15,7 @@ dt_eval <- function(x) {
 # even when data.table isn't attached
 dt_funs <- c(
   "copy", "dcast", "melt", "nafill",
-  "fcase", "fcoalesce", "fintersect", "frank", "frankv", "fsetdiff", "funion",
+  "fcase", "fcoalesce", "fifelse", "fintersect", "frank", "frankv", "fsetdiff", "funion",
   "setcolorder", "setnames"
 )
 add_dt_wrappers <- function(env) {
@@ -129,8 +129,9 @@ dt_squash_call <- function(x, env, data, j = TRUE) {
     quote(.GRP)
   } else if (is_call(x, "cur_group_rows")) {
     quote(.I)
-  } else if (is_call(x, "if_else")) {
+  } else if (is_call(x, "if_else") || is_call(x, "ifelse")) {
     x[[1L]] <- quote(fifelse)
+    x[-1] <- lapply(x[-1], dt_squash, env, data, j = j)
     x
   } else if (is_call(x, "n", n = 0)) {
     quote(.N)
diff --git a/man/relocate.dtplyr_step.Rd b/man/relocate.dtplyr_step.Rd
index 0f74bc3a7..f1ad62020 100644
--- a/man/relocate.dtplyr_step.Rd
+++ b/man/relocate.dtplyr_step.Rd
@@ -1,5 +1,5 @@
 % Generated by roxygen2: do not edit by hand
-% Please edit documentation in R/step-subset-relocate.R
+% Please edit documentation in R/step-colorder-relocate.R
 \name{relocate.dtplyr_step}
 \alias{relocate.dtplyr_step}
 \title{Relocate variables using their names}
diff --git a/tests/testthat/test-tidyeval.R b/tests/testthat/test-tidyeval.R
index 6a2c35a14..ba2e070c3 100644
--- a/tests/testthat/test-tidyeval.R
+++ b/tests/testthat/test-tidyeval.R
@@ -183,15 +183,29 @@ test_that("scoped verbs produce nice output", {
     dt %>% summarise_all(~ n()) %>% show_query(),
     expr(DT[, .(x = .N)])
   )
+})
+
+test_that("mask if_else/ifelse & coalesce with data.table versions", {
+  dt <- lazy_dt(data.table(x = 1:5), "DT")
 
-  # mask if_else & coalesce with data.table versions, #112
+  # if_else/ifelse & coalesce to fifelse & fcoalesce, #112
+  expect_equal(
+    dt %>% mutate(y = if_else(x < 3, 1, 2)) %>% show_query(),
+    expr(copy(DT)[, `:=`(y = fifelse(x < 3, 1, 2))])
+  )
+  expect_equal(
+    dt %>% mutate(y = ifelse(x < 3, 1, 2)) %>% show_query(),
+    expr(copy(DT)[, `:=`(y = fifelse(x < 3, 1, 2))])
+  )
   expect_equal(
-    dt %>% summarise_all(~if_else(. > 0, -1, 1)) %>% show_query(),
-    expr(DT[ , .(x = fifelse(x > 0, -1, 1))])
+    dt %>% mutate(y = coalesce(x, 1)) %>% show_query(),
+    expr(copy(DT)[, `:=`(y = fcoalesce(x, 1))])
   )
+
+  # tidyeval works inside if_else, #220
   expect_equal(
-    dt %>% summarise_all(~coalesce(., 1)) %>% show_query(),
-    expr(DT[ , .(x = fcoalesce(x, 1))])
+    dt %>% mutate(y = if_else(.data$x < 3, 1, 2)) %>% show_query(),
+    expr(copy(DT)[, `:=`(y = fifelse(x < 3, 1, 2))])
   )
 })