lolplyr: Adapting dplyr for lists-of-lists

#' Add new variables
#'
#' @param .data  a list.
#'
#' @param \dots  name-value pairs of expressions.
#'
#'               \code{mutate_elem} and \code{transmute_elem} work with data
#'               taking as input one element at a time, they can be used to
#'               transform individual elements of the list.
#'               \code{mutate} and \code{transmute} work in a similar manner
#'               as their dplyr counterparts, i.e. they are applied to whole
#'               lists/vectors.
#'
#'               \emph{Notice:} in \code{mutate} or \code{transmute} expressions
#'               need to evaluate to vectors or lists of the same length as
#'               \code{length(.data)}, otherwise the functions will throw an error.
#'
#' @param gather before calculating the summaries, each variable is first passed through
#'               the \code{gather} function. Common choices for \code{gather} are
#'               \code{\link{unlist}}, or \code{\link{unlist_with_nas}}. If data is to be
#'               taken as-is, use \code{\link{identity}} function.
#'
#' @param warn   if \code{FALSE} (default) it ignores the warnings throwed by
#'               purrr's \code{\link[purrr]{transpose}}.
#'
#' @seealso \code{\link[dplyr]{mutate}}, \code{\link[dplyr]{transmute}}
#'
#' @examples
#'
#' lst <- list(
#'   list(a = 1:3),
#'   list(a = 5),
#'   list(a = c(NA, 4)),
#'   list(a = NA),
#'   list(a = 100)
#' )
#'
#' lst %>% mutate_elem(norm_a = a/sum(a, na.rm = TRUE))
#'
#' @export

mutate_elem <- function(.data, ...) UseMethod("mutate_elem")

#' @export

mutate_elem.list <- function(.data, ...) {
  dots <- quos(...)
  lapply(.data, function(.x) {
    eval_dots <- try_eval(dots, data = as.list(.x))
    .x[names(eval_dots)] <- eval_dots
    .x
  })
}


#' @rdname mutate_elem
#' @export

transmute_elem <- function(.data, ...) UseMethod("transmute_elem")

#' @export

transmute_elem.list <- function(.data, ...) {
  dots <- quos(...)
  lapply(.data, function(.x) try_eval(dots, data = as.list(.x)))
}



# dplyr-like mutate


mutate_list_impl <- function(.data, ..., gather = unlist_with_nas,
                             warn = FALSE, drop = FALSE) {

  if (!is_simple_list(.data))
    stop(".data is not a list")

  dots <- quos(...)

  if (!warn)
    transp <- function(x) suppressWarnings(transpose(x))
  else
    transp <- transpose

  t_data <- transp(.data)
  t_data <- lapply(t_data, gather)

  eval_dots <- try_eval(dots, data = t_data)

  eval_dots <- lapply(eval_dots, function(.x) {
    .x <- as.list(.x)
    if (length(.x) != length(.data))
      stop("mutate has changed the dimensions of data")
    names(.x) <- names(.data)
    .x
  })

  if (!drop) {
    out <- t_data
    out[names(eval_dots)] <- eval_dots
  } else {
    out <- eval_dots
  }

  transp(out)
}


#' @examples
#'
#' as_lol(mtcars) %>%
#'   transmute(
#'     avg_mpg = rep_list(mean(mpg), length(mpg)),
#'     cyl = paste("it has", cyl, "cylinders")
#'   ) %>% .[1:5]
#'
#' as_lol(mtcars) %>%
#'   mutate(
#'     avg_mpg = rep_list(mean(mpg), length(mpg)),
#'     cyl = cyl + 1000
#'   ) %>% .[1:5]
#'
#' list(
#'   a = list(x = 1),
#'   b = list(x = 2),
#'   c = list(x = 3),
#'   list(x = 4)
#'  ) %>%
#'   mutate(y = x / sum(x))
#'
#' @rdname mutate_elem
#' @export

mutate.list <- function(.data, ..., gather = unlist_with_nas, warn = FALSE) {

  mutate_list_impl(.data, ..., gather = gather,
                   warn = warn, drop = FALSE)
}

#' @rdname mutate_elem
#' @export

transmute.list <- function(.data, ..., gather = unlist_with_nas, warn = FALSE) {

  mutate_list_impl(.data, ..., gather = gather,
                   warn = warn, drop = TRUE)
}