Implement the 'cache.vars' chunk option; some style changes

R/knitr-cache.R

@@ -16,6 +16,13 @@
 #'
 #' @export
 cache_eng_python <- (function() {
+  closure <- environment()
+  dill <- NULL
+
+  cache_path <- function(path) {
+    paste(path, "pkl", sep=".")
+  }
+
   check_cache_available <- function(options) {
     MINIMUM_PYTHON_VERSION <- "3.7"
     MINIMUM_DILL_VERSION <- "0.3.6"
@@ -29,10 +36,11 @@ cache_eng_python <- (function() {
     }
 
     # is the module 'dill' loadable and recent enough?
-    dill <- tryCatch(import("dill"), error = identity)
+    closure$dill <- tryCatch(import("dill"), error = identity)
     if (!inherits(dill, "error")) {
       dill_version <- as_numeric_version(dill$`__version__`)
       if (dill_version >= MINIMUM_DILL_VERSION)
+        cache_initialize()
         return(TRUE)
     } else {
       # handle non-import error
@@ -47,16 +55,15 @@ cache_eng_python <- (function() {
   }
 
   cache_available <- function(options) {
-    available <- knitr::opts_knit$get("reticulate.cache")
-    if (is.null(available)) {
-      available <- check_cache_available(options)
-      knitr::opts_knit$set(reticulate.cache = available)
-    }
-    available
+    if (is.null(closure$.cache_available))
+      closure$.cache_available <- check_cache_available(options)
+    .cache_available
   }
 
-  cache_path <- function(path) {
-    paste(path, "pkl", sep=".")
+  cache_initialize <- function() {
+    # save imported objects by reference when possible
+    dill.session <- import("dill.session")
+    dill.session[["settings"]][["refimported"]] <- TRUE
   }
 
   cache_exists <- function(options) {
@@ -65,23 +72,31 @@ cache_eng_python <- (function() {
 
   cache_load <- function(options) {
     if (!cache_available(options)) return()
-    dill <- import("dill")
     dill$load_module(filename = cache_path(options$hash), module = "__main__")
   }
 
-  filter <- NULL
   r_obj_filter <- function() {
-    if (is.null(filter)) {
-      filter <<- py_eval("lambda obj: obj.name == 'r' and type(obj.value) is __builtins__.__R__")
+    if (is.null(closure$.r_obj_filter)) {
+      expr <- "lambda obj: obj.name == 'r' and type(obj.value) is __builtins__.__R__"
+      closure$.r_obj_filter <- py_eval(expr)
     }
-    filter
+    .r_obj_filter
   }
 
   cache_save <- function(options) {
     if (!cache_available(options)) return()
-    dill <- import("dill")
+
+    # when only inclusion filters are specified, it works as an allowlist
+    if (!is.null(options$cache.vars)) {
+      exclude <- NULL  # the R object won't be saved unless specified by cache.vars
+      include <- options$cache.vars
+    } else {
+      exclude <- r_obj_filter()
+      include <- NULL
+    }
+
     tryCatch({
-      dill$dump_module(cache_path(options$hash), refimported = TRUE, exclude = r_obj_filter())
+      dill$dump_module(cache_path(options$hash), exclude = exclude, include = include)
     }, error = function(e) {
       cache_purge(options$hash)
       stop(e)

tests/testthat/resources/eng-reticulate-cache.Rmd

@@ -18,3 +18,21 @@ open('py_chunk_executed', 'w').close()
 ```{python, cache = FALSE}
 print(x + 1)
 ```
+
+The `cache.vars` chunk option may be used to select only a subset of variables
+from the global environment to be cached.
+
+```{python cache-vars, cache.vars = 'x'}
+x = 42
+spam = "Lovely SPAM! Wonderful SPAM!"
+```
+
+In a second excution of the document, `spam` should not be present, because it wasn't saved.
+
+```{python, cache = FALSE}
+def print_globals():
+    for name, value in globals().items():
+        if not name.startswith('__'):
+            print(name, ": ", value, sep="")
+print_globals()
+```

tests/testthat/test-python-knitr-cache.R

@@ -9,11 +9,12 @@ test_that("An R Markdown document using reticulate can be rendered with cache fe
   flag_file <- "py_chunk_executed"
   rmd_prefix <- "eng-reticulate-cache"
   rmd_file <- paste(rmd_prefix, "Rmd", sep=".")
+  cache_path <- paste(rmd_prefix, "cache", sep="_")
 
   withr::with_dir("resources", local({
     withr::defer({
       unlink(flag_file)
-      unlink(paste(rmd_prefix, "cache", sep="_"), recursive = TRUE)
+      unlink(cache_path, recursive = TRUE)
     })
 
     # cache file is created
@@ -26,5 +27,13 @@ test_that("An R Markdown document using reticulate can be rendered with cache fe
     output <- rmarkdown::render(rmd_file, quiet = TRUE)
     expect_false(file.exists(flag_file))
     expect_true(file.exists(output))
+
+    # the 'spam' variable should not be cached in the 'cache-vars' block
+    main <- import_main()
+    dill <- import("dill")
+    py_del_attr(main, "spam")
+    session_file <- Sys.glob(paste0(cache_path, "/*/cache-vars_*.pkl"))
+    dill$load_module(session_file)
+    expect_false("spam" %in% names(main))
   }))
 })