From 26156fafedf18f23013bb7ff79b9a033db080e7c Mon Sep 17 00:00:00 2001
From: Romain Francois <romain@rstudio.com>
Date: Thu, 18 Feb 2021 18:19:58 +0100
Subject: [PATCH] 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
---
 R/group-by.r                   | 11 ++---------
 tests/testthat/test-group-by.r |  3 +++
 2 files changed, 5 insertions(+), 9 deletions(-)

diff --git a/R/group-by.r b/R/group-by.r
index 24994b14d2..43612aaa61 100644
--- a/R/group-by.r
+++ b/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}()`."),
diff --git a/tests/testthat/test-group-by.r b/tests/testthat/test-group-by.r
index fe74ebfe2a..5135f662f3 100644
--- a/tests/testthat/test-group-by.r
+++ b/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 ------------------------------------------------------------------