Create a function that can handle both list or symbol as input argument

[noveld]

I've created a function, select_or_return , that should handle both column symbols and lists as input arguments. When providing it with a column symbol, it works as expected. However, when I provide a list, I encounter an error related to the ensym() function from the rlang package.

How could I create a function that can handle both a list or symbol as input argument?

library(rlang)
library(dplyr)
#> 
#> Attaching package: 'dplyr'
#> The following objects are masked from 'package:stats':
#> 
#>     filter, lag
#> The following objects are masked from 'package:base':
#> 
#>     intersect, setdiff, setequal, union

select_or_return <- function(data, x) {
  
  x_sym <- ensym(x)
  
  # If x is a symbol
  if (is_symbol(x_sym)) {
    return(data %>% select(!!x_sym))
  } 
  # If x is a list
  else if (is.list(x)) {
    return(x)
  } 
  # For other cases
  else {
    stop("x must be a symbol representing a column or a list.")
  }
}

# Example usage:

# Create a sample dataframe
df <- data.frame(a = 1:5, b = 6:10)

# Use the function with a column name as a symbol
print(select_or_return(df, a))
#>   a
#> 1 1
#> 2 2
#> 3 3
#> 4 4
#> 5 5

# Use the function with a list
print(select_or_return(df, list(1,2,3)))
#> Error in `ensym()`:
#> ! Can't convert to a symbol.
#> Backtrace:
#>     ▆
#>  1. ├─base::print(select_or_return(df, list(1, 2, 3)))
#>  2. ├─global select_or_return(df, list(1, 2, 3))
#>  3. │ └─rlang::ensym(x)
#>  4. └─rlang::abort(message = message)

^{Created on 2023-08-10 with reprex v2.0.2}

[nirgrahamuk]

here is one way

select_or_return <- function(data, x) {
  x_sym <- enexpr(x)
  if(inherits(x_sym,"call")){
    force(x)
    x_sym <- x
  }

  if (is.list(x_sym)) {
    return(x_sym)
  } 

  if (is_symbol(x_sym)) {
    return(data %>% select(!!x_sym))
  } 
    stop("x must be a symbol representing a column or a list.")
}

[noveld]

Thanks a lot! Why does the function return an error I use it in a wrapper function? How can I use it in both contexts?

select_from_symbol <- function(data, x) {
  
  x_expr <- rlang::enexpr(x)
  
  if(inherits(x_expr, "call")) {
    x_val <- eval(x_expr)
    
    if (is.character(x_val)) {
      return(x_val)
    }
  } else if (rlang::is_symbol(x_expr)) {
    column_name <- rlang::as_string(x_expr)
    
    # Check if column name exists in the data
    if (!is.null(column_name) && !(column_name %in% colnames(data))) {
      stop(paste("Column", column_name, "not found in data"))
    }
    
    return(data[[column_name]])
  }
  stop("x must be a symbol representing a column or a character vector.")
}

print(select_from_symbol(mtcars, mpg))
#>  [1] 21.0 21.0 22.8 21.4 18.7 18.1 14.3 24.4 22.8 19.2 17.8 16.4 17.3 15.2 10.4
#> [16] 10.4 14.7 32.4 30.4 33.9 21.5 15.5 15.2 13.3 19.2 27.3 26.0 30.4 15.8 19.7
#> [31] 15.0 21.4
print(select_from_symbol(mtcars, as.character(mtcars$mpg)))
#>  [1] "21"   "21"   "22.8" "21.4" "18.7" "18.1" "14.3" "24.4" "22.8" "19.2"
#> [11] "17.8" "16.4" "17.3" "15.2" "10.4" "10.4" "14.7" "32.4" "30.4" "33.9"
#> [21] "21.5" "15.5" "15.2" "13.3" "19.2" "27.3" "26"   "30.4" "15.8" "19.7"
#> [31] "15"   "21.4"

wrapper_function <- function(data, x){
  data <- select_from_symbol(data,x)
  return(data)
}

print(wrapper_function(mtcars, mpg))
#> Error in select_from_symbol(data, x): Column x not found in data
print(wrapper_function(mtcars, as.character(mtcars$mpg)))
#> Error in select_from_symbol(data, x): Column x not found in data

^{Created on 2023-08-10 with reprex v2.0.2}

[nirgrahamuk]

It seems to me your wrapper has the same requirements and challenge as what it tries to wrap; so it would need to have essentially the same code within it.

[noveld]

I don't get why it is working when I use the function on it's own but as soon as I wrap it inside another one it is not. I guess it has something to do with the environment.