R/purrr-misc.R

In Nmisc: Miscellaneous Functions Used at 'Numeract LLC'

#' Keep or discard elements
#' 
#' @description
#' 
#' \code{keep_at()} keeps only the elements from specific positions
#'   while \code{discard_at()} does the opposite.
#'   The functions are wrappers around \code{purrr::keep} and
#'   \code{purrr::discard}, respectively.
#' 
#' @param .x A list or a vector.
#' @param .at A character vector (names), a numeric vector (positions),
#'   a symbol or or a list generated by 
#'   \code{\link[tidyselect:select_helpers]{tidyselect}} 
#'   select helpers.
#' 
#' @return A list or a vector.
#' 
#' @examples
#' x <- c("First" = 1, "Second" = 2, "Last" = 3)
#' keep_at(x, "Second")
#' keep_at(x, Second)
#' keep_at(x, 2)
#' keep_at(x, starts_with("Sec"))
#' #> Second 
#' #>      2
#' 
#' keep_at(x, ends_with("t"))
#' #> First  Last 
#' #>     1     3
#' 
#' x <- c(1, 2, 3)
#' discard_at(x, 1)
#' #> Second   Last 
#' #>      2      3
#' 
#' @seealso \code{\link[purrr:keep]{purrr::keep}}
#' 
#' @export
keep_at <- function(.x, .at) {
    
    vars <- names(.x)
    if (length(vars) == 0) {
        # tidyselect not available 
        sel <- .at
    } else {
        quo <- rlang::enquo(.at)
        helpers <- tidyselect::vars_select_helpers
        quo <- rlang::env_bury(quo, !!! helpers)
        sel <- tidyselect::vars_select(vars, !! quo)
    }
    
    .p <- probe_at(.x, sel)
    purrr::keep(.x, .p)
}


#' @export
#' @rdname keep_at
discard_at <- function(.x, .at) {
    
    vars <- names(.x)
    if (length(vars) == 0) {
        # tidyselect not available 
        sel <- .at
    } else {
        quo <- rlang::enquo(.at)
        helpers <- tidyselect::vars_select_helpers
        quo <- rlang::env_bury(quo, !!! helpers)
        sel <- tidyselect::vars_select(vars, !! quo)
    }
    
    .p <- probe_at(.x, sel)
    purrr::discard(.x, .p)
}


probe_at <- function(.x, .at) {
    
    if (anyNA(.at)) stop("`.at` must not contain any NA's")
    
    if (is.character(.at)) {
        names(.x) %in% .at
    } else if (is.numeric(.at)) {
        seq_along(.x) %in% as.integer(.at)
    } else {
        stop("`.at` must be character (names) or numeric (positions)")
    }
}