grantmcdermott / etwfe

Extended two-way fixed effects
https://grantmcdermott.com/etwfe/
Other
50 stars 11 forks source link

Support NSE? #1

Closed grantmcdermott closed 1 year ago

grantmcdermott commented 1 year ago

Something like the following, adapted from subset, internally. Would add a few microseconds in overhead, but it handles both expressions and strings.

nse_func = function(vars, data) {
  nl = as.list(seq_along(data))
  names(nl) = names(data)
  vars = eval(substitute(vars), nl, parent.frame())
  if (is.numeric(vars)) vars = names(data)[vars]
  # Simple return object to demonstrate
  head(data[, vars, drop = FALSE], 2)
}

# Handles both strings and expressions 
nse_func("mpg", mtcars)
#>               mpg
#> Mazda RX4      21
#> Mazda RX4 Wag  21
nse_func(mpg, mtcars)
#>               mpg
#> Mazda RX4      21
#> Mazda RX4 Wag  21

# Vectors are fine too (although not expected for etwfe)
nse_func(c("mpg", "cyl"), mtcars)
#>               mpg cyl
#> Mazda RX4      21   6
#> Mazda RX4 Wag  21   6
nse_func(c(mpg, cyl), mtcars)
#>               mpg cyl
#> Mazda RX4      21   6
#> Mazda RX4 Wag  21   6
nse_func(c(mpg:cyl, hp), mtcars)
#>               mpg cyl  hp
#> Mazda RX4      21   6 110
#> Mazda RX4 Wag  21   6 110
nse_func(names(mtcars), mtcars)
#>               mpg cyl disp  hp drat    wt  qsec vs am gear carb
#> Mazda RX4      21   6  160 110  3.9 2.620 16.46  0  1    4    4
#> Mazda RX4 Wag  21   6  160 110  3.9 2.875 17.02  0  1    4    4

Created on 2022-10-15 with reprex v2.0.2