R/add_margins.R

In kimjam/qrcutils: Contains utility functions for data pipeline project.

#' @title add_margins
#' @description Function to add column and row totals to a dataframe. Assumes character columns
#' followed by all numeric columns.
#'
#' @param df Dataframe to add margins to.
#' @param rowsum Boolean to indicate if you want row sums in a Total column. Defaults to TRUE.
#'
#' @return Returns dataframe with margins added.
#' @export
add_margins <- function(df, rowsum = TRUE) {

  if (!tail(vapply(df, is.numeric, logical(1)), 1)) {
    stop('Expects character columns followed by all numeric columns.')
  }
  if ('grouped_df' %in% class(df)) {
    df %<>% dplyr::ungroup()
  }

  colclasses <- vapply(X = df,
                       FUN = is.numeric,
                       FUN.VALUE = logical(1))

  numstart <- match(TRUE, colclasses)

  rowmargin <- vector('list', ncol(df))
  for (i in seq(numstart - 1)) {
    if (i < (numstart - 1)) {
      rowmargin[[i]] <- ''
    } else {
      rowmargin[[i]] <- 'Total'
    }
  }

  for (i in seq(numstart, ncol(df))) {
    rowmargin[[i]] <- sum(df[, i])
  }

  df %<>%
    rbind(., rowmargin) %>%
    setNames(replace(names(df), names(df) == '', 'None')) %>%
    {if (rowsum) {
      dplyr::mutate(., Total = rowSums(.[numstart:ncol(.)]))
    } else {.}}

  return(df)
}