Problem: Can I augment an existing function with checking code for its return value?

I’m still experimenting with function body modification, and I’ve downscaled my ambitions and thought about what a type check helper should look like.

Yesterday’s post showed that a simple function could be used to add checks to an existing function.

This post is about writing a function to add a check for the return value of a function.

Example function to be checked

calc_value <- function(a=1L, b=2, c=3) {
  (a + b) / c
}

check_return_value()

The function for adding checks to another function is shown below.

It operates by:

• taking a function as its first argument
• all subsequent arguments are interpreted as boolean statements for checking the return value with the name res
• create a code block combining all these tests
• create a complete call to the function with the result captured in res
• create a new function that calls the original function, then checks the return value for any errors before returning it
• return the augmented version of the function

#-----------------------------------------------------------------------------
#' add checks for the return value of a function
#'
#' @param fun existing function passed in a symbol
#' @param ... list of checks to add for the result of the function assigned to `res`
#'
#' @return new function (with the same function signature as `fun`) with tests
#'         added for the return value
#-----------------------------------------------------------------------------
check_return_value <- function(fun, ...) {
  
  # Capture all the tests and turn each one into a stopifnot() call
  checks <- rlang::exprs(...)
  for (i in seq(checks)) {
    checks[[i]] <- bquote(stopifnot(isTRUE(.(checks[[i]]))))
  }
  
  # Bind all these checks into a single block
  checks <- rlang::call2('{', splice(checks))
  
  # Create a call to the function
  fun_sym     <- rlang::enexpr(fun)
  call_to_fun <- rlang::call2(fun_sym, splice(syms(names(formals(fun)))))
  
  # Concatentate the call to the function with the checks
  new_body <- bquote({
    res <- .(call_to_fun)
    .(checks)
    return(res)
  })
  
  # create and return the new function
  rlang::new_function(args = formals(fun), body = new_body)
}

calc_value_checked_return <- check_return_value(calc_value, !is_na(res))

calc_value_checked_return

## function (a = 1L, b = 2, c = 3) 
## {
##     res <- calc_value(a, b, c)
##     {
##         stopifnot(isTRUE(!is_na(res)))
##     }
##     return(res)
## }
## <environment: 0x7f8695305340>

> calc_value_checked_return()
[1] 1

> calc_value_checked_return(b = NA)  # the result of the calculation will be NA
Error: isTRUE(!is_na(res)) is not TRUE

> calc_value_checked_return(c = 0) # this will cause a divide by zero error -> Inf
[1] Inf

check_args()

As per yesterday’s post, a function can also have its arguments checked. Yesterday’s add_checks() function has been renamed to check_args() and repeated here.

#-----------------------------------------------------------------------------
#' add checks to an existing function and return a new function
#'
#' @param fun existing function passed in a symbol
#' @param ... list of checks to add in front of function body
#'
#' @return new function (with the same function signature as `fun`) and the same
#'         body as `fun` with a block of tests inserted at the start of the function
#-----------------------------------------------------------------------------
check_args <- function(fun, ...) {
  
  # Capture all the tests and turn each one into a stopifnot() call
  checks <- rlang::exprs(...)
  for (i in seq(checks)) {
    checks[[i]] <- bquote(stopifnot(isTRUE(.(checks[[i]]))))
  }
  
  # Bind all these checks into a single block
  checks <- rlang::call2('{', splice(checks))
  
  # Concatentate the test block with the original body
  new_body <- bquote({
    .(checks)
    .(body(fun))
  })
  
  # create and return the new function
  rlang::new_function(args = formals(fun), body = new_body)
}

Now add a check to ensure that the argument c is never 0. This is in addition to the checks for the return value of the function as a whole.

calc_value_checked_args_and_return <- check_args(calc_value_checked_return, c != 0)
calc_value_checked_args_and_return

## function (a = 1L, b = 2, c = 3) 
## {
##     {
##         stopifnot(isTRUE(c != 0))
##     }
##     {
##         res <- calc_value(a, b, c)
##         {
##             stopifnot(isTRUE(!is_na(res)))
##         }
##         return(res)
##     }
## }
## <environment: 0x7f869531ee00>

> calc_value_checked_args_and_return()
[1] 1

> calc_value_checked_args_and_return(b = NA) # An error should be raised because of an NA return value
Error: isTRUE(!is_na(res)) is not TRUE

> calc_value_checked_args_and_return(c = 0) # An error should be raised because `c` should never be 0
Error: isTRUE(c != 0) is not TRUE

Conclusion

• A combination of check_args() and check_return_value() means that assertions about the arguments to a function and the return value from a function can be checked for types, and these checks can be added to the original function to produce a new function
• These functions aren’t currently pipe-friendly because of the way the fun argument is used and the way piping with %>% uses . for argument representation. Ideally the functions should work like:

new_func <- calc_value %>% 
  check_return_value(!is.na(res)) %>% 
  check_args(c != 0)