R/group_by.R

In twolodzko/lolplyr: Adapting dplyr for lists-of-lists

#' Group by one or more variables
#'
#' \code{group_by} re-arranges list, by creating a hierarchical
#' list groupped according to variables listed in \dots. To
#' back-transform the groupped list, use \code{unlist}.
#'
#' @param .data a list.
#' @param \dots variables to group by.
#' @param x     an object.
#' @param warn  throw warnings when encountering errors in selecting variables
#'              (\code{FALSE} by default).
#'
#' @export

group_by.list <- function(.data, ..., warn = TRUE) {

  dots <- quos(...)

  if (length(dots) == 0L)
    return(.data)

  out <- list()
  dict <- list()

  walk(.data, function(.x) {
    vars <- try_select_vars(names(.x), dots, warn = warn)
    hsh <- hash(.x[vars])
    dict[[hsh]] <<- .x[vars]
    out[[hsh]] <<- c(out[[hsh]], list(.x))
  })

  structure(
    out,
    class = c("grouped_list", "list"),
    orig_names = names(.data),
    orig_hashes = hash_by(.data),
    dict = dict
  )
}

#' @export

ungroup.list <- function(x, ...) x

#' @rdname group_by.list
#' @export

ungroup.grouped_list <- function(x, ...) {
  out <- list()
  walk(x, function(.x) {
    out <<- c(out, .x)
  })
  new_hashes <- hash_by(out)
  idx <- match(attr(x, "orig_hashes"), new_hashes)
  out <- out[idx]
  names(out) <- attr(x, "orig_names")
  out
}

#' @rdname group_by.list
#' @export

is_grouped_list <- function(x) {
  isTRUE(
    is_simple_list(x) &&
    "grouped_list" %in% class(x) &&
    "dict" %in% names(attributes(x)) &&
    "orig_hashes" %in% names(attributes(x))
  )
}

#' @export

summarise.grouped_list <- function(.data, ...,
                                   gather = unlist_with_nas,
                                   warn = FALSE) {

  out <- lapply(.data, function(.x) {
    summarise(.x, ..., gather = gather, warn = warn)
  })

  unname(map2(out, attr(.data, "dict")[names(out)], ~ c(.y, .x)))
}