How do I check whether a custom statistic passed to gtsummary's add_stat() is on the label or level location?

Question

I'm currently building a custom function to evaluate pairwise differences between groups in tbl_summary/tbl_svysummary. When using %>% add_stat(fns = everything() ~ custom_stat, location = ~"level", the function should return a df where nrow is equal to the number of levels of a given categorical variable, while location = ~"label" should return a single row. I'm not picking up on the location argument being passed to custom_stat, so how do I get that information in order to format the output accordingly?

I have searched through the gtsummary documentation and other requests/questions but have been unable to find an answer.

---

REPREX:

Let's take the following code:

library(tidyverse)
library(gtsummary)

print_one <- function(data, variable, by, tbl, loc = "label", ...) {
  if (loc == "label") {
    return(tibble(Res = 1))
  } else {
    if (class(data[[variable]]) != "character" & class(data[[variable]]) != "factor" ) {
      return(tibble(Res = c()))
    } else {
      tibble(Res = 1:length(unique(data[[variable]])))
    }
  }
}

trial %>% 
  tbl_summary(by = trt, include = c(age, stage)) %>% 
  add_stat(fns = everything() ~ print_one, location = ~ "level")

I want a way to forward the gtsummary location argument (i.e., either the value "label" or "level" into the print_one function, so that when location = ~ "label" I get the following output:

and when location = ~ "level" it should be:

So my question is: How do I get the "loc" argument in my print_one function to receive the value of the "location" argument of gtsummary's add_stat function?

EDIT: Add-on question: Is there a way to manually specify additional arguments (i.e., the "loc" argument) to the print_one function when I call it in %>% add_stat(fns = everything() ~ print_one)?

Answer 1

Okay, so after prodding around more I came to the conclusion that intercepting the value of the location argument by the print_one() function was not possible.

Instead, I came up with a different way to apply the function to allow users to specify the location:

library(tidyverse)
library(gtsummary)

print_one <- function(data, variable, location) {
  if (location == "label") {
    return(tibble(Res = 1))
  } else {
    if (class(data[[variable]]) != "character" & class(data[[variable]]) != "factor" ) {
      return(tibble(Res = c()))
    } else {
      tibble(Res = 1:length(unique(data[[variable]])))
    }
  }
}

apply_print_one <- function(tbl, location = "label") {
  fun <- function(data, variable, by, tbl, ...) {
    print_one(data, variable, location)
  }
  tbl %>% add_stat(fns = everything() ~ fun, location = ~ location)
}

trial %>% 
  tbl_summary(by = trt, include = c(age, stage)) %>% 
  apply_print_one("level")

That way, calling apply_print_one() creates a temporary function fun which calls print_one() with the right argument values and supplies the same location parameter to add_stat() all at once - while making the code look neater to users.