Fix deferred_run() in knitr

NEWS.md

@@ -3,6 +3,8 @@
 * `deferred_run()` now reports the number of executed expressions with
   a message.
 
+* `deferred_run()` can now be run at any point in a knitr file (#235).
+
 ,* `local_tempfile()` now writes `lines` in UTF-8 (#210) and always uses 
   `\n` for newlines (#216).
 

R/defer-exit.R

@@ -1,43 +1,12 @@
-#' Find frame for exit handlers
-#' @noRd
-#'
-#' @description
-#' - If knitr is running and `defer()` is run at top-level in a chunk,
-#'   `exit_frame()` returns the first frame on the stack that inherits
-#'   from the knitr namespace. This way handlers are run when knitr
-#'   returns.
-#'
-#' - If `getOption("withr.hook_source")` is `TRUE`, and `defer()` is
-#'   run at top-level in a script being evaluated by [base::source()],
-#'   `exit_frame()` returns the frame of `source()`. This way handlers
-#'   are run when the script has finished running.
-#'
-exit_frame <- function(envir,
-                       frames = as.list(sys.frames()),
-                       calls = as.list(sys.calls())) {
+# Find first knitr frame on the stack
+knitr_exit_frame <- function(envir) {
+  frames <- as.list(sys.frames())
+
   frame_loc <- frame_loc(envir, frames)
   if (!frame_loc) {
     return(envir)
   }
 
-  if (in_knitr(envir)) {
-    out <- knitr_frame(envir, frames, calls, frame_loc)
-    if (!is.null(out)) {
-      return(out)
-    }
-  }
-
-  if (isTRUE(getOption("withr.hook_source"))) {
-    out <- source_frame(envir, frames, calls, frame_loc)
-    if (!is.null(out)) {
-      return(out)
-    }
-  }
-
-  envir
-}
-
-knitr_frame <- function(envir, frames, calls, frame_loc) {
   knitr_ns <- asNamespace("knitr")
 
   # This doesn't handle correctly the recursive case (knitr called
@@ -52,11 +21,17 @@ knitr_frame <- function(envir, frames, calls, frame_loc) {
   NULL
 }
 
-source_frame <- function(envir, frames, calls, frame_loc) {
-  i <- frame_loc
+source_exit_frame <- function(envir) {
+  frames <- as.list(sys.frames())
+  calls <- as.list(sys.calls())
+
+  i <- frame_loc(envir, frames)
+  if (!i) {
+    return(envir)
+  }
 
   if (i < 4) {
-    return(NULL)
+    return(envir)
   }
 
   is_call <- function(x, fn) {
@@ -65,16 +40,16 @@ source_frame <- function(envir, frames, calls, frame_loc) {
   calls <- as.list(calls)
 
   if (!is_call(calls[[i - 3]], quote(source))) {
-    return(NULL)
+    return(envir)
   }
   if (!is_call(calls[[i - 2]], quote(withVisible))) {
-    return(NULL)
+    return(envir)
   }
   if (!is_call(calls[[i - 1]], quote(eval))) {
-    return(NULL)
+    return(envir)
   }
   if (!is_call(calls[[i - 0]], quote(eval))) {
-    return(NULL)
+    return(envir)
   }
 
   frames[[i - 3]]

R/defer.R

@@ -76,18 +76,19 @@ defer <- function(expr, envir = parent.frame(), priority = c("first", "last")) {
   }
 
   priority <- match.arg(priority, choices = c("first", "last"))
-  after <- priority == "last"
 
-  thunk <- as.call(list(function() expr))
+  if (knitr_in_progress() && identical(envir, knitr::knit_global())) {
+    return(defer_knitr(expr, envir, priority = priority))
+  }
 
-  # Don't handle `source()` and `knit()` specially by default
-  # to avoid a performance hit
-  hook_source <- getOption("withr.hook_source")
-  hook_knitr <- getOption("knitr.in.progress")
-  if (!is.null(hook_source) || !is.null(hook_knitr)) {
-    envir <- exit_frame(envir)
+  # Don't handle `source()` by default to avoid a performance hit
+  if (!is.null(getOption("withr.hook_source"))) {
+    envir <- source_exit_frame(envir)
   }
 
+  thunk <- as.call(list(function() expr))
+  after <- priority == "last"
+
   do.call(
     base::on.exit,
     list(thunk, TRUE, after),
@@ -108,15 +109,19 @@ defer_parent <- function(expr, priority = c("first", "last")) {
 #' @rdname defer
 #' @export
 deferred_run <- function(envir = parent.frame()) {
-  if (knitr_in_progress()) {
-    stop("Can't run `deferred_run()` in a knitted document")
-  }
-  if (is_top_level_global_env(envir)) {
-    handlers <- the$global_exits
+  if (knitr_in_progress() && identical(envir, knitr::knit_global())) {
+    # The handlers are thunks so we don't need to clear them.
+    # They will only be run once.
+    frame <- knitr_exit_frame(envir)
+    handlers <- knitr_handlers(frame)
   } else {
-    handlers <- frame_exits(envir)
+    if (is_top_level_global_env(envir)) {
+      handlers <- the$global_exits
+    } else {
+      handlers <- frame_exits(envir)
+    }
+    deferred_clear(envir)
   }
-  deferred_clear(envir)
 
   n <- length(handlers)
   i <- 0L
@@ -220,6 +225,51 @@ is_top_level_global_env <- function(envir, frames = sys.frames()) {
 }
 
 
+# This picks up knitr's first frame on the stack and registers the
+# handler there. To avoid mixing up knitr's own exit handlers with
+# ours, we don't hook directly but instead save the list of handlers
+# as an attribute on the frame environment. This allows `deferred_run()`
+# to run our handlers without running the knitr ones.
+defer_knitr <- function(expr, envir, priority = c("first", "last")) {
+  priority <- match.arg(priority, choices = c("first", "last"))
+
+  envir <- knitr_exit_frame(envir)
+  handler <- as.call(list(function() expr))
+
+  handlers <- knitr_handlers(envir)
+
+  # Add `on.exit` hook if run for first time
+  if (!length(handlers)) {
+    defer_knitr_run(envir)
+  }
+
+  if (priority == "first") {
+    handlers <- c(list(handler), handlers)
+  } else {
+    handlers <- c(handlers, list(handler))
+  }
+  attr(envir, "withr_knitr_handlers") <- handlers
+
+  invisible(NULL)
+}
+
+knitr_handlers <- function(envir) {
+  attr(envir, "withr_knitr_handlers") %||% list()
+}
+
+# Evaluate `handlers` lazily so we get the latest version
+defer_knitr_run <- function(
+  envir,
+  handlers = knitr_handlers(envir)
+) {
+  defer(envir = envir, {
+    for (expr in handlers) {
+      eval(expr, envir)
+    }
+  })
+}
+
+
 # Augment rlang with withr features such as knitr support
 on_load({
   on_package_load("rlang", local({

R/utils.R

@@ -93,3 +93,7 @@ is_interactive <- function() {
   }
   interactive()
 }
+
+`%||%` <- function(x, y) {
+  if (is.null(x)) y else x
+}

tests/testthat/_snaps/defer.md

@@ -28,3 +28,53 @@
     Message
       Ran 1/3 deferred expressions
 
+# can't run `deferred_run()` in knitr
+
+    Code
+      writeLines(readLines(out))
+    Output
+
+
+      ```r
+      withr::deferred_run()
+      ```
+
+      ```
+      ## No deferred expressions to run
+      ```
+
+      ```r
+      defer(writeLines('1'))
+      writeLines('2')
+      ```
+
+      ```
+      ## 2
+      ```
+
+      ```r
+      defer(writeLines('3'))
+      ```
+
+      ```r
+      writeLines('4')
+      ```
+
+      ```
+      ## 4
+      ```
+
+      ```r
+      withr::deferred_run()
+      ```
+
+      ```
+      ## 3
+      ## 1
+      ```
+
+      ```
+      ## Ran 2/2 deferred expressions
+      ```
+
+

tests/testthat/test-defer.R

@@ -33,13 +33,25 @@ test_that("can't run `deferred_run()` in knitr", {
     skip_if_cannot_knit()
 
     rmd <- local_tempfile(fileext = ".Rmd")
+    out <- local_tempfile(fileext = ".md")
     writeLines(rmd, text = "
 ```{r}
 withr::deferred_run()
 ```
+```{r}
+defer(writeLines('1'))
+writeLines('2')
+defer(writeLines('3'))
+```
+```{r}
+writeLines('4')
+withr::deferred_run()
+```
 ")
-    expect_error(
-      suppressMessages(rmarkdown::render(rmd, quiet = TRUE)),
-      "in a knitted document"
-    )
+
+    knitr::knit(rmd, out, quiet = TRUE)
+
+    expect_snapshot({
+      writeLines(readLines(out))
+    })
 })