Home > Enterprise >  Function argument as value or column name for data.table
Function argument as value or column name for data.table

Time:10-29

I want my function to be able to take a value or a column name. How can I do this with data.table?

library(data.table)

df <- data.table(a = c(1:5), 
                 b = c(5:1), 
                 c = c(1, 3, 5, 3, 1))

myfunc <- function(val) {
  df[a >= val]
}

# This works:
myfunc(2)

# This does not work: 
myfunc("c")

If I define my function as:

myfunc <- function(val) {
  df[a >= get(val)]
}

# This doesn't work:
myfunc(2)

# This works: 
myfunc("c")

What is the best way to resolve this?

EDIT: Thanks all for the responses, I think I like dww's answer the best. I wish it was as easy as in dplyr, where I can do:

myfunc <- function(val) {
  df %>%
    filter(a >= {{val}})
}

# Both work:
myfunc(2)
myfunc(c)

CodePudding user response:

If you build and parse the whole expression, then you can evaluate it in its entirety. For example

myfunc <- function(val) {
  df[eval(parse(text=paste("a >= ", val)))]
}

Though relying on a function that lets you mix values and variable names in the same parameter might be dangerous. Especially in the case where you actually wanted to match on character values rather than variable names. If you passed in the whole expression you could do

myfunc <- function(expr) {
  expr <- substitute(expr)
  df[eval(expr)]
}

myfunc(a>=3)
myfunc(a>=c)

CodePudding user response:

The question did not actually define the desired behavior so we assume that df must be a data.table and if a character string is passed then the column of that name should be returned and if a number is passed then those rows whose a column exceed that number should be returned.

Define an S3 generic and methods for character and default.

myfunc <- function(x, data = df) UseMethod("myfunc")
myfunc.character <- function(x, data = df) data[[x]]
myfunc.default <- function(x, data = df) data[a > x]
 
myfunc(2)
##    a b c
## 1: 3 3 5
## 2: 4 2 3
## 3: 5 1 1

myfunc("c")
## [1] 1 3 5 3 1
  • Related