R/03_mutate_ifx.R

In moodymudskipper/tidyx: modified tidyverse functions for interractive use

#' generalized mutate_if / mutate_at / mutate_all
#'
#' takes the features of \code{mutate_if} and \code{mutate_at}, in case we want to combine
#' both, and supports a formula or function for the .at parameter, which will
#' be applied on names. Also supports grouping through \code{.by}.
#'
#' @inheritParams dplyr::mutate_if
#' @param .at equivalent of .vars in mutate_at, but enhanced
#' @param .if equivalent of .predicate in mutate_if
#' @param .by a vector or list of variables as characters or given by \code{vars}
#' @param .keep if \code{"new"}, \code{transmute} will be called instead of \code{mutate}
#'
#' @export
#'
#' @examples
#' iris %>% head %>% mutate_ifx(funs(z = mean),.at = ~startsWith(.,"Pet"))
#' iris %>% head %>% mutate_ifx(funs(z = mean),.at = vars(Petal.Length))
#' iris %>% head %>% mutate_ifx(funs(z = mean),.at = vars(Petal.Length), .keep = "new")
mutate_ifx <- function(.tbl, .funs, ..., .at = names(.tbl), .if = ~TRUE,
                       .by = NULL, .rowwise=FALSE, .keep = c("all","new"))
{
  # checks
  if(!is.null(.by) && .rowwise)
    stop("You can't use rowwise and .by at the same time")
  .keep <- match.arg(.keep)
  .by <- dplyr:::tbl_at_syms(.data, .by)

  # convert at to logical and fuse with .if
  .at <- switch(mode(.at),
                list = names(.tbl) %in% dplyr:::tbl_at_vars(.tbl, .at),
                logical = .at,
                numeric = `[<-`(rep(FALSE,length(.tbl)), .at, value = TRUE),
                character = names(.tbl) %in% .at,
                call = mapx(names(.tbl), .at, .t = "lgl"),
                `function` = mapx(names(.tbl), .at, .t = "lgl"))
  .if        <- if (is.logical(.if)) .if else  mapx(.tbl, .if, .t = "lgl")
  .predicate <- .if & .at
  funs <- dplyr:::manip_if(.tbl, .predicate, .funs,
                           rlang::enquo(.funs), rlang::caller_env(), ...)

  # add default .i and .n arguments to functions if not conflicting
  for (i in seq_along(funs)) {
    .f <- funs[[i]][[2]][[1]]
    if (is.name(.f)) .f <- eval.parent(.f)
    if (!".i" %in% formalArgs(.f))
      formals(.f) <- c(formals(.f), list(.i = i))
    else message(".f already has a .i formal, the `.i` ",
                 "feature will be disabled for this function")
    if (!".n" %in% formalArgs(.f))
      formals(.f) <- c(formals(.f),list(.n = rlang::names2(funs)[i]))
    else message(".f already has a .n formal, the `.n` ",
                 "feature will be disabled for this function")
    funs[[i]][[2]][[1]] <- .f
  }

  # group if relevant, then mutate/transmute and ungroup
  if (!length(.by) && ! .rowwise)
    if (.keep == "all") dplyr::mutate(.tbl, !!!funs)
    else dplyr::transmute(.tbl, !!!funs)
  else {
    groups <- dplyr::group_vars(.tbl)
    x <- if (.rowwise) dplyr::rowwise(.tbl) else dplyr::group_by(.tbl,!!!.by, add = TRUE)
    x <- if (.keep == "all") dplyr::mutate(x, !!!funs) else dplyr::transmute(x, !!!funs)
    x <- dplyr::group_by(ungroup(x), !!!rlang::syms(groups))
    x
  }
}