Modify the body of a function programmatically

modify_expr.R

# 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"