dplyr at home

cols.R

#' Column selector for `[` and `[[`
#'
#' @param ... Column names provided as `character` or `symbol` (could be
#'   extended to allow other inputs!)
#' @return A `numeric` vector of column indices to select
#' 
#' @examples
#' mtcars[cols(cyl, wt)] |> head()
#' 
#' mtcars[cols(mpg:wt)] |> head()
#' 
#' mtcars[-cols(cyl, wt)] |> head()
#' 
#' mtcars[cols(cyl, wt, -wt)] |> head()
#' 
#' penguins[cols(starts_with("bill"))] |> head()
#' 
#' 
cols <- function(...) {
  indexing_frame <- NA_integer_
  indexing_call <- NULL

  # find indexing op call
  for (frame in seq_len(sys.nframe())) {
    call <- sys.call(-frame)
    if (is.call(call) && as.character(call[[1L]]) %in% c("[", "[[")) {
      indexing_frame <- -frame
      indexing_call <- call
      break
    }
  }

  # throw error if we can't find our parent call
  if (is.null(indexing_call)) {
    stop("`cols` can only be used within `[` and `[[` indexing operators")
  }

  # find our "context" object
  ctx <- eval(indexing_call[[2L]], envir = sys.frame(-indexing_frame))
  tidyselect::eval_select(rlang::expr(c(...)), data = ctx)
}