From 21a3791a5173190ad885abf968952993ca2e37ce Mon Sep 17 00:00:00 2001
From: olivroy <olivierroy71@hotmail.com>
Date: Thu, 26 Sep 2024 15:25:53 -0400
Subject: [PATCH] fixup cols_add() in corner cases

---
 NEWS.md                           |  2 ++
 R/cols_add.R                      | 25 +++++++++++--------------
 tests/testthat/_snaps/cols_add.md | 25 +++++++++++++++++++++++++
 tests/testthat/test-cols_add.R    | 30 ++++++++++++++++++++++++++++--
 4 files changed, 66 insertions(+), 16 deletions(-)
 create mode 100644 tests/testthat/_snaps/cols_add.md

diff --git a/NEWS.md b/NEWS.md
index 38a719a69a..48f37ef67b 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -63,6 +63,8 @@
 
 * `vec_fmt_*()` (and incidentally `cols_nanoplot()`) should be faster now (@olivroy, #1888, #1891).
 
+* `cols_add()` works in more cases (#1893).
+
 # gt 0.11.0
 
 ## New features
diff --git a/R/cols_add.R b/R/cols_add.R
index 4f69ae7450..43da0939e1 100644
--- a/R/cols_add.R
+++ b/R/cols_add.R
@@ -384,10 +384,6 @@ cols_add <- function(
     resolved_column_before <- NULL
   } else if (length(resolved_column_before) != 1) {
 
-    if (length(resolved_column_before) < 1) {
-      cli::cli_abort("The expression used for `.before` did not resolve a column.")
-    }
-
     if (length(resolved_column_before) > 1) {
       cli::cli_abort("The expression used for `.before` resolved multiple columns.")
     }
@@ -404,14 +400,15 @@ cols_add <- function(
     resolved_column_after <- NULL
   } else if (length(resolved_column_after) != 1) {
 
-    if (length(resolved_column_after) < 1) {
-      cli::cli_abort("The expression used for `.after` did not resolve a column.")
-    }
-
     if (length(resolved_column_after) > 1) {
       cli::cli_abort("The expression used for `.after` resolved multiple columns.")
     }
   }
+  
+  if (length(resolved_column_after) == 1 && resolved_column_after == colnames(data_tbl)[NCOL(data_tbl)]) {
+    # if requesting the last column to add after, use NULL instead.
+    resolved_column_after <- NULL
+  } 
 
   # Stop function if expressions are given to both `.before` and `.after`
   if (!is.null(resolved_column_before) && !is.null(resolved_column_after)) {
@@ -455,9 +452,9 @@ cols_add <- function(
     } else {
       updated_data_tbl <-
         dplyr::bind_cols(
-          dplyr::select(data_tbl, 1:(before_colnum - 1)),
+          dplyr::select(data_tbl, 1:(dplyr::all_of(before_colnum) - 1)),
           data_tbl_new_cols,
-          dplyr::select(data_tbl, before_colnum:ncol(data_tbl))
+          dplyr::select(data_tbl, dplyr::all_of(before_colnum):ncol(data_tbl))
         )
     }
 
@@ -473,7 +470,7 @@ cols_add <- function(
 
   } else if (is.null(resolved_column_before) && !is.null(resolved_column_after)) {
 
-    if (resolved_column_after == nrow(data_tbl)) {
+    if (resolved_column_after == ncol(data_tbl)) {
 
       updated_data_tbl <-
         vctrs::vec_cbind(
@@ -493,14 +490,14 @@ cols_add <- function(
 
       if (after_colnum >= ncol(data_tbl)) {
         updated_data_tbl <-
-          dplyr::bind_cols(
+          vctrs::vec_cbind(
             data_tbl,
             data_tbl_new_cols
           )
       } else {
         updated_data_tbl <-
           dplyr::bind_cols(
-            dplyr::select(data_tbl, 1:after_colnum),
+            dplyr::select(data_tbl, 1:dplyr::all_of(after_colnum)),
             data_tbl_new_cols,
             dplyr::select(data_tbl, (after_colnum + 1):ncol(data_tbl))
           )
@@ -508,7 +505,7 @@ cols_add <- function(
 
 
       after_colnum <- which(boxh_df[["var"]] == resolved_column_after)
-
+      
       updated_boxh_df <-
         dplyr::bind_rows(
           boxh_df[1:after_colnum, ],
diff --git a/tests/testthat/_snaps/cols_add.md b/tests/testthat/_snaps/cols_add.md
new file mode 100644
index 0000000000..f7a6f0c729
--- /dev/null
+++ b/tests/testthat/_snaps/cols_add.md
@@ -0,0 +1,25 @@
+# cols_add() errors with bad input
+
+    Code
+      cols_add(tbl, x = 1, .after = 2, .before = 3)
+    Condition
+      Error in `cols_add()`:
+      ! Expressions cannot be given to both `.before` and `.after`.
+    Code
+      cols_add(tbl, x = 1, .after = 15)
+    Condition
+      Error in `vars_select_eval()`:
+      ! Can't select columns past the end.
+      i Location 15 doesn't exist.
+      i There are only 9 columns.
+    Code
+      cols_add(tbl, x = 1, .before = c(1, 2))
+    Condition
+      Error in `cols_add()`:
+      ! The expression used for `.before` resolved multiple columns.
+    Code
+      cols_add(tbl, x = 1, .after = c(1, 2))
+    Condition
+      Error in `cols_add()`:
+      ! The expression used for `.after` resolved multiple columns.
+
diff --git a/tests/testthat/test-cols_add.R b/tests/testthat/test-cols_add.R
index 2b26c831aa..2fa95d107e 100644
--- a/tests/testthat/test-cols_add.R
+++ b/tests/testthat/test-cols_add.R
@@ -1,8 +1,34 @@
 test_that("cols_add() works", {
   tbl <- gt(exibble)
   expect_no_error(
-    cols_add(tbl, x = 1, .before = 3)
+    dat <- cols_add(tbl, x = 1, .before = 1)
+  )
+  expect_equal(dat$`_boxhead`$var[1], "x")
+  
+  expect_no_error(
+    dat <- cols_add(tbl, x = 1, .before = 2)
+  )
+  expect_equal(dat$`_boxhead`$var[2], "x")
+  expect_no_error(
+    dat <- cols_add(tbl, x = 1, .after = 1)
+  )
+  expect_equal(dat$`_boxhead`$var[2], "x")
+  expect_no_error(
+    dat <- cols_add(tbl, x = 1, .after = dplyr::last_col())
+  )
+})
+
+test_that("cols_add() errors with bad input", {
+  tbl <- gt(exibble)
+  
+  expect_snapshot(
+    error = TRUE, {
+      cols_add(tbl, x = 1, .after = 2, .before = 3)
+      cols_add(tbl, x = 1, .after = 15)
+      cols_add(tbl, x = 1, .before = c(1, 2))
+      cols_add(tbl, x = 1, .after = c(1, 2))
+      
+    }
   )
-  expect_equal(cols_add(), what)
 })