Skip to content

Instantly share code, notes, and snippets.

@wch
Last active January 13, 2023 21:30
Embed
What would you like to do?
Modify the body of a function programmatically
# This function recurses into `expr`, looking for `search_expr`. If it finds
# it, it will return `replace_expr`. Notably, `replace_expr` can contain rlang
# operators like `!!`, and perform rlang substitution on them.
modify_expr <- function(
expr,
search_expr,
replace_expr
) {
if (typeof(expr) != "language") {
stop("modify_expr only works on language objects (AKA quoted expressions)!")
}
# Look for search_expr in the code. If it matches, return replace_expr, but do
# substitution on it.
if (identical(
removeSource(search_expr),
removeSource(expr[seq_along(search_expr)])
)) {
result <- rlang::inject(rlang::expr(!!replace_expr))
return(result)
}
# If we get here, then expr is a language object, but not the one we were
# looking for. Recurse into it. We would use lapply here if we could, but that
# returns a list instead of language object. So we'll iterate using a for loop
# instead.
for (i in seq_along(expr)) {
# Recurse only if expr[[i]] is a language object.
if (typeof(expr[[i]]) == "language") {
expr[[i]] <- modify_expr(expr[[i]], search_expr, replace_expr)
}
}
expr
}
# ============================================================================
# Example usage
# ============================================================================
# This is the function we want to modify.
f <- function(x) {
if (x < 2) {
return("x is less than 2")
} else {
return("x is 2 or larger")
}
}
# Do the modification. We want to prepend an `if` condition to the `if-else`
# chain.
new_body <- modify_expr(
body(f),
quote(
if (x < 2) {
return("x is less than 2")
}
),
quote(
if (x < 1) {
return("x is less than 1")
} else !!expr
)
)
# Save the body back
body(f) <- new_body
# The modified version of `f` has a `if (x < 1)` condition at the beginning.
f
#> function (x)
#> {
#> if (x < 1) {
#> return("x is less than 1")
#> }
#> else if (x < 2) {
#> return("x is less than 2")
#> }
#> else {
#> return("x is 2 or larger")
#> }
#> }
f(0)
#> [1] "x is less than 1"
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment