From db2830124f0ad27896a8563e8f18c7cc96a89e3e Mon Sep 17 00:00:00 2001
From: "Matthijs S. Berends" <m.s.berends@umcg.nl>
Date: Wed, 15 Feb 2023 19:48:34 +0100
Subject: [PATCH] fix for using `dplyr::select()`

---
 DESCRIPTION             |  2 +-
 NEWS.md                 |  2 +-
 R/aa_helper_functions.R | 46 ++++++++++++++++++++++++-----------------
 3 files changed, 29 insertions(+), 21 deletions(-)

diff --git a/DESCRIPTION b/DESCRIPTION
index 11795c7c..9a8f6cee 100644
--- a/DESCRIPTION
+++ b/DESCRIPTION
@@ -1,5 +1,5 @@
 Package: AMR
-Version: 1.8.2.9130
+Version: 1.8.2.9131
 Date: 2023-02-15
 Title: Antimicrobial Resistance Data Analysis
 Description: Functions to simplify and standardise antimicrobial resistance (AMR)
diff --git a/NEWS.md b/NEWS.md
index 7bb3f380..01fc31c1 100755
--- a/NEWS.md
+++ b/NEWS.md
@@ -1,4 +1,4 @@
-# AMR 1.8.2.9130
+# AMR 1.8.2.9131
 
 *(this beta version will eventually become v2.0! We're happy to reach a new major milestone soon!)*
 
diff --git a/R/aa_helper_functions.R b/R/aa_helper_functions.R
index dcb5178f..3042c948 100755
--- a/R/aa_helper_functions.R
+++ b/R/aa_helper_functions.R
@@ -890,11 +890,12 @@ get_current_data <- function(arg_name, call) {
     !is.null(x) && is.data.frame(x)
   }
 
-  # try a manual (base R) method, by going over all underlying environments with sys.frames()
-  for (env in sys.frames()) {
-
-    # dplyr support ----
-    if (!is.null(env$mask) && is.function(env$mask$current_rows) && (valid_df(env$data) || valid_df(env$`.data`))) {
+  frms <- sys.frames()
+  
+  # check dplyr environments to support dplyr groups
+  with_mask <- vapply(FUN.VALUE = logical(1), frms, function(e) !is.null(e$mask))
+  for (env in frms[which(with_mask)]) {
+    if (is.function(env$mask$current_rows) && (valid_df(env$data) || valid_df(env$`.data`))) {
       # an element `.data` or `data` (containing all data) and `mask` (containing functions) will be in the environment when using dplyr verbs
       # we use their mask$current_rows() to get the group rows, since dplyr::cur_data_all() is deprecated and will be removed in the future
       # e.g. for `example_isolates %>% group_by(ward) %>% mutate(first = first_isolate(.))`
@@ -907,21 +908,28 @@ get_current_data <- function(arg_name, call) {
       }
       rows <- tryCatch(env$mask$current_rows(), error = function(e) seq_len(NROW(df)))
       return(df[rows, , drop = FALSE])
-      
-      # base R support ----
-    } else if (!is.null(env$`.Generic`)) {
-      # don't check `".Generic" %in% names(env)`, because in R < 3.2, `names(env)` is always NULL
+    }
+  }
 
-      if (valid_df(env$xx)) {
-        # an element `xx` will be in the environment for rows + cols in base R, e.g. `example_isolates[c(1:3), carbapenems()]`
-        return(env$xx)
-      } else if (valid_df(env$x)) {
-        # an element `x` will be in the environment for only cols in base R, e.g. `example_isolates[, carbapenems()]`
-        return(env$x)
-      }
-      
-      # scoped dplyr support ----
-    } else if (!is.null(names(env)) && all(c(".tbl", ".vars", ".cols") %in% names(env), na.rm = TRUE) && valid_df(env$`.tbl`)) {
+  # now go over all underlying environments looking for other dplyr and base R selection environments
+  with_generic <- vapply(FUN.VALUE = logical(1), frms, function(e) !is.null(e$`.Generic`))
+  for (env in frms[which(with_generic)]) {
+    if (valid_df(env$`.data`)) {
+      # an element `.data` will be in the environment when using dplyr::select()
+      return(env$`.data`)
+    } else if (valid_df(env$xx)) {
+      # an element `xx` will be in the environment for rows + cols in base R, e.g. `example_isolates[c(1:3), carbapenems()]`
+      return(env$xx)
+    } else if (valid_df(env$x)) {
+      # an element `x` will be in the environment for only cols in base R, e.g. `example_isolates[, carbapenems()]`
+      return(env$x)
+    }
+  }
+  
+  # now a special case for dplyr's 'scoped' variants
+  with_tbl <- vapply(FUN.VALUE = logical(1), frms, function(e) valid_df(e$`.tbl`))
+  for (env in frms[which(with_tbl)]) {
+    if (!is.null(names(env)) && all(c(".tbl", ".vars", ".cols") %in% names(env), na.rm = TRUE)) {
       # an element `.tbl` will be in the environment when using scoped dplyr variants, with or without `dplyr::vars()`
       # (e.g. `dplyr::summarise_at()` or `dplyr::mutate_at()`)
       return(env$`.tbl`)