From 9baa22be43af54053d5586323b901951faa4a9a5 Mon Sep 17 00:00:00 2001
From: Daniel <mail@danielluedecke.de>
Date: Thu, 21 Nov 2024 10:28:04 +0100
Subject: [PATCH] Allow `n()` in `data_modify()` (#535)

* Allow `n()` in `data_modify()`

* lintr, styler

* Update NEWS.md

Co-authored-by: Etienne Bacher <52219252+etiennebacher@users.noreply.github.com>

* Update R/data_modify.R

Co-authored-by: Etienne Bacher <52219252+etiennebacher@users.noreply.github.com>

* comments

* fix test

* update rd

* modify error msg

* error on invalid function

* move news item

---------

Co-authored-by: Etienne Bacher <52219252+etiennebacher@users.noreply.github.com>
---
 NEWS.md                           |  3 ++
 R/data_modify.R                   | 21 +++++++++++--
 man/data_modify.Rd                |  6 +++-
 tests/testthat/test-data_modify.R | 52 +++++++++++++++++++++++++++++++
 4 files changed, 79 insertions(+), 3 deletions(-)

diff --git a/NEWS.md b/NEWS.md
index 8fc8a29ca..663efa310 100644
--- a/NEWS.md
+++ b/NEWS.md
@@ -19,6 +19,9 @@ CHANGES
 * `data_read()` no longer shows warning about forthcoming breaking changes
   in upstream packages when reading `.RData` files.
 
+* `data_modify()` now recognizes `n()`, for example to create an index for data groups
+  with `1:n()` (#535).
+
 BUG FIXES
 
 * `describe_distribution()` no longer errors if the sample was too sparse to compute
diff --git a/R/data_modify.R b/R/data_modify.R
index e7744c1f5..3e30b8f68 100644
--- a/R/data_modify.R
+++ b/R/data_modify.R
@@ -22,6 +22,9 @@
 #'   character vector is provided, you may not add further elements to `...`.
 #' - Using `NULL` as right-hand side removes a variable from the data frame.
 #'   Example: `Petal.Width = NULL`.
+#' - For data frames (including grouped ones), the function `n()` can be used to count the
+#'   number of observations and thereby, for instance, create index values by
+#'   using `id = 1:n()` or `id = 3:(n()+2)` and similar.
 #'
 #' Note that newly created variables can be used in subsequent expressions,
 #' including `.at` or `.if`. See also 'Examples'.
@@ -92,7 +95,8 @@
 #'   grouped_efc,
 #'   c12hour_c = center(c12hour),
 #'   c12hour_z = c12hour_c / sd(c12hour, na.rm = TRUE),
-#'   c12hour_z2 = standardize(c12hour)
+#'   c12hour_z2 = standardize(c12hour),
+#'   id = 1:n()
 #' )
 #' head(new_efc)
 #'
@@ -145,6 +149,11 @@ data_modify.default <- function(data, ...) {
 data_modify.data.frame <- function(data, ..., .if = NULL, .at = NULL, .modify = NULL) {
   dots <- eval(substitute(alist(...)))
 
+  # error for data frames with no rows...
+  if (nrow(data) == 0) {
+    insight::format_error("`data` is an empty data frame. `data_modify()` only works for data frames with at least one row.") # nolint
+  }
+
   # check if we have dots, or only at/modify ----
 
   if (length(dots)) {
@@ -201,6 +210,10 @@ data_modify.grouped_df <- function(data, ..., .if = NULL, .at = NULL, .modify =
   # the data.frame method later...
   dots <- match.call(expand.dots = FALSE)[["..."]]
 
+  # error for data frames with no rows...
+  if (nrow(data) == 0) {
+    insight::format_error("`data` is an empty data frame. `data_modify()` only works for data frames with at least one row.") # nolint
+  }
 
   grps <- attr(data, "groups", exact = TRUE)
   grps <- grps[[".rows"]]
@@ -352,8 +365,12 @@ data_modify.grouped_df <- function(data, ..., .if = NULL, .at = NULL, .modify =
   # finally, we can evaluate expression and get values for new variables
   symbol_string <- insight::safe_deparse(symbol)
   if (!is.null(symbol_string) && all(symbol_string == "n()")) {
-    # "special" functions
+    # "special" functions - using "n()" just returns number of rows
     new_variable <- nrow(data)
+  } else if (!is.null(symbol_string) && length(symbol_string) == 1 && grepl("\\bn\\(\\)", symbol_string)) {
+    # "special" functions, like "1:n()" or similar - but not "1:fun()"
+    symbol_string <- str2lang(gsub("n()", "nrow(data)", symbol_string, fixed = TRUE))
+    new_variable <- try(with(data, eval(symbol_string)), silent = TRUE)
   } else {
     # default evaluation of expression
     new_variable <- try(with(data, eval(symbol)), silent = TRUE)
diff --git a/man/data_modify.Rd b/man/data_modify.Rd
index 042962e03..28533ecea 100644
--- a/man/data_modify.Rd
+++ b/man/data_modify.Rd
@@ -30,6 +30,9 @@ type of expression cannot be mixed with other expressions, i.e. if a
 character vector is provided, you may not add further elements to \code{...}.
 \item Using \code{NULL} as right-hand side removes a variable from the data frame.
 Example: \code{Petal.Width = NULL}.
+\item For data frames (including grouped ones), the function \code{n()} can be used to count the
+number of observations and thereby, for instance, create index values by
+using \code{id = 1:n()} or \code{id = 3:(n()+2)} and similar.
 }
 
 Note that newly created variables can be used in subsequent expressions,
@@ -109,7 +112,8 @@ new_efc <- data_modify(
   grouped_efc,
   c12hour_c = center(c12hour),
   c12hour_z = c12hour_c / sd(c12hour, na.rm = TRUE),
-  c12hour_z2 = standardize(c12hour)
+  c12hour_z2 = standardize(c12hour),
+  id = 1:n()
 )
 head(new_efc)
 
diff --git a/tests/testthat/test-data_modify.R b/tests/testthat/test-data_modify.R
index 9bb0a92d6..a7a153c43 100644
--- a/tests/testthat/test-data_modify.R
+++ b/tests/testthat/test-data_modify.R
@@ -353,6 +353,16 @@ test_that("data_modify errors for non df", {
 })
 
 
+test_that("data_modify errors for empty data frames", {
+  data(mtcars)
+  x <- mtcars[1, ]
+  expect_error(
+    data_modify(x[-1, ], new_var = 5),
+    regex = "empty data frame"
+  )
+})
+
+
 test_that("data_modify errors for non df", {
   data(efc)
   a <- "center(c22hour)" # <---------------- error in variable name
@@ -492,6 +502,20 @@ test_that("data_modify works with functions that return character vectors", {
 })
 
 
+test_that("data_modify 1:n() and similar works in (grouped) data frames", {
+  data(mtcars)
+  out <- data_modify(mtcars, Trials = 1:n()) # nolint
+  expect_identical(out$Trials, 1:32)
+  x <- data_group(mtcars, "gear")
+  out <- data_modify(x, Trials = 1:n()) # nolint
+  expect_identical(out$Trials[out$gear == 3], 1:15)
+  expect_identical(out$Trials[out$gear == 4], 1:12)
+  out <- data_modify(x, Trials = 3:(n() + 2))
+  expect_identical(out$Trials[out$gear == 3], 3:17)
+  expect_identical(out$Trials[out$gear == 4], 3:14)
+})
+
+
 test_that("data_modify .if/.at arguments", {
   data(iris)
   d <- iris[1:5, ]
@@ -550,3 +574,31 @@ test_that("data_modify .if/.at arguments", {
   out <- data_modify(d, new_length = Petal.Length * 2, .if = is.numeric, .modify = round)
   expect_equal(out$new_length, c(3, 3, 3, 3, 3), ignore_attr = TRUE)
 })
+
+
+skip_if_not_installed("withr")
+
+withr::with_environment(
+  new.env(),
+  test_that("data_modify 1:n() and similar works in (grouped) data frames inside function calls", {
+    data(mtcars)
+    x <- data_group(mtcars, "gear")
+
+    foo <- function(d) {
+      out <- data_modify(d, Trials = 1:n()) # nolint
+      out$Trials
+    }
+    expect_identical(
+      foo(x),
+      c(
+        1L, 2L, 3L, 1L, 2L, 3L, 4L, 4L, 5L, 6L, 7L, 5L, 6L, 7L, 8L,
+        9L, 10L, 8L, 9L, 10L, 11L, 12L, 13L, 14L, 15L, 11L, 1L, 2L, 3L,
+        4L, 5L, 12L
+      )
+    )
+  })
+)
+
+test_that("data_modify errors on non-defined function", {
+  expect_error(data_modify(iris, Species = foo()))
+})