1#' Apply a function to list-elements of a list
2#'
3#' `lmap()`, `lmap_at()` and `lmap_if()` are similar to
4#' `map()`, `map_at()` and `map_if()`, with the
5#' difference that they operate exclusively on functions that take
6#' \emph{and} return a list (or data frame). Thus, instead of mapping
7#' the elements of a list (as in \code{.x[[i]]}), they apply a
8#' function `.f` to each subset of size 1 of that list (as in
9#' `.x[i]`). We call those elements `list-elements`).
10#'
11#' Mapping the list-elements `.x[i]` has several advantages. It
12#' makes it possible to work with functions that exclusively take a
13#' list or data frame. It enables `.f` to access the attributes
14#' of the encapsulating list, like the name of the components it
15#' receives. It also enables `.f` to return a larger list than
16#' the list-element of size 1 it got as input. Conversely, `.f`
17#' can also return empty lists. In these cases, the output list is
18#' reshaped with a different size than the input list `.x`.
19#' @param .x A list or data frame.
20#' @param .f A function that takes and returns a list or data frame.
21#' @inheritParams map_if
22#' @inheritParams map_at
23#' @inheritParams map
24#' @return If `.x` is a list, a list. If `.x` is a data
25#'   frame, a data frame.
26#' @family map variants
27#' @export
28#' @examples
29#' # Let's write a function that returns a larger list or an empty list
30#' # depending on some condition. This function also uses the names
31#' # metadata available in the attributes of the list-element
32#' maybe_rep <- function(x) {
33#'   n <- rpois(1, 2)
34#'   out <- rep_len(x, n)
35#'   if (length(out) > 0) {
36#'     names(out) <- paste0(names(x), seq_len(n))
37#'   }
38#'   out
39#' }
40#'
41#' # The output size varies each time we map f()
42#' x <- list(a = 1:4, b = letters[5:7], c = 8:9, d = letters[10])
43#' x %>% lmap(maybe_rep)
44#'
45#' # We can apply f() on a selected subset of x
46#' x %>% lmap_at(c("a", "d"), maybe_rep)
47#'
48#' # Or only where a condition is satisfied
49#' x %>% lmap_if(is.character, maybe_rep)
50#'
51#'
52#' # A more realistic example would be a function that takes discrete
53#' # variables in a dataset and turns them into disjunctive tables, a
54#' # form that is amenable to fitting some types of models.
55#'
56#' # A disjunctive table contains only 0 and 1 but has as many columns
57#' # as unique values in the original variable. Ideally, we want to
58#' # combine the names of each level with the name of the discrete
59#' # variable in order to identify them. Given these requirements, it
60#' # makes sense to have a function that takes a data frame of size 1
61#' # and returns a data frame of variable size.
62#' disjoin <- function(x, sep = "_") {
63#'   name <- names(x)
64#'   x <- as.factor(x[[1]])
65#'
66#'   out <- lapply(levels(x), function(level) {
67#'     as.numeric(x == level)
68#'   })
69#'
70#'   names(out) <- paste(name, levels(x), sep = sep)
71#'   out
72#' }
73#'
74#' # Now, we are ready to map disjoin() on each categorical variable of a
75#' # data frame:
76#' iris %>% lmap_if(is.factor, disjoin)
77#' mtcars %>% lmap_at(c("cyl", "vs", "am"), disjoin)
78lmap <- function(.x, .f, ...) {
79  lmap_at(.x, seq_along(.x), .f, ...)
80}
81
82#' @rdname lmap
83#' @export
84lmap_if <- function(.x, .p, .f, ..., .else = NULL) {
85  sel <- probe(.x, .p)
86
87  .x <- lmap_at(.x, which(sel), .f, ...)
88
89  if (!is_null(.else)) {
90    .x <- lmap_at(.x, which(!sel), .else, ...)
91  }
92
93  .x
94}
95
96#' @rdname lmap
97#' @export
98lmap_at <- function(.x, .at, .f, ...) {
99  if (is_formula(.f)) {
100    .f <- as_mapper(.f, ...)
101  }
102
103  where <- at_selection(names(.x), .at)
104  sel <- inv_which(.x, where)
105
106  out <- vector("list", length(.x))
107  for (i in seq_along(.x)) {
108    res <-
109      if (sel[[i]]) {
110        .f(.x[i], ...)
111      } else {
112        .x[i]
113      }
114    stopifnot(is.list(res))
115    out[[i]] <- res
116  }
117
118  maybe_as_data_frame(flatten(out), .x)
119}
120