Move the ungroup() special handling down to add_coomputed_columsn()

related to https://github.com/tidyverse/dplyr/pull/5767/files/bbe2e5be39946ed4f7bea5149d84e2612133fc51#diff-1ec33e14b938e61e20e0bef01b54aa084a603bfcbf3449c84535293f0f6ca9a5 comment on #5767

Author: romainfrancois

R/group-by.r

@@ -149,15 +149,8 @@ group_by_prepare <- function(.data, ..., .add = FALSE, .dots = deprecated(), add
   }
 
   # If any calls, use mutate to add new columns, then group by those
-  ungrouped <- ungroup(.data)
-  atts <- attributes(.data)
-  atts <- append(
-    attributes(ungrouped),
-    atts[setdiff(names(atts), c("names", "row.names", "groups", "class"))]
-  )
-  attributes(ungrouped) <- atts
+  computed_columns <- add_computed_columns(.data, new_groups, "group_by")
 
-  computed_columns <- add_computed_columns(ungrouped, new_groups, "group_by")
   out <- computed_columns$data
   group_names <- computed_columns$added_names
 
@@ -188,7 +181,7 @@ add_computed_columns <- function(.data, vars, .fn = "group_by") {
     # TODO: use less of a hack
     if (inherits(.data, "data.frame")) {
       cols <- withCallingHandlers(
-        mutate_cols(.data, !!!vars),
+        mutate_cols(ungroup(.data), !!!vars),
         error = function(e) {
           abort(c(
             glue("Problem adding computed columns in `{.fn}()`."),

tests/testthat/test-group-by.r

@@ -538,6 +538,9 @@ test_that("group_by() keeps attributes unrelated to the grouping (#5760)", {
 
   d3 <- group_by(d2, y, .add = TRUE)
   expect_equal(attr(d2, "foo"), "bar")
+
+  d4 <- group_by(d2, y2 = y * 2, .add = TRUE)
+  expect_equal(attr(d2, "foo"), "bar")
 })
 
 # Errors ------------------------------------------------------------------