From 6e74c1e2a0b6ceec3e734ec1dea850998e5df07f Mon Sep 17 00:00:00 2001
From: Daniel Falbel <dfalbel@gmail.com>
Date: Thu, 20 Jul 2023 14:22:07 -0300
Subject: [PATCH] Fix knitr output remapping from a clean new session (#1418)

---
 R/output.R                                          |  1 +
 inst/python/rpytools/output.py                      |  6 +++++-
 .../_snaps/python-knitr-engine/knitr-print2.md      |  4 ++++
 tests/testthat/test-python-knitr-engine.R           | 13 ++++++++++++-
 4 files changed, 22 insertions(+), 2 deletions(-)
 create mode 100644 tests/testthat/_snaps/python-knitr-engine/knitr-print2.md

diff --git a/R/output.R b/R/output.R
index 64ea98dc9..f48712486 100644
--- a/R/output.R
+++ b/R/output.R
@@ -36,6 +36,7 @@ set_knitr_python_stdout_hook <- function() {
     # if knitr is already in progress here, this means that python was initialized
     # during a chunk execution. We have to force an instant remap as the hook won't
     # have a chance to run for that chunk.
+    # In such cases `context.__enter__` is never called.
     if (isTRUE(getOption('knitr.in.progress'))) set_output_streams(tty = FALSE)
   } else {
     setHook(
diff --git a/inst/python/rpytools/output.py b/inst/python/rpytools/output.py
index bbe4da41f..ef4c83b59 100644
--- a/inst/python/rpytools/output.py
+++ b/inst/python/rpytools/output.py
@@ -153,14 +153,18 @@ def __init__(self, r_stdout, r_stderr, tty):
     self.r_stdout = r_stdout
     self.r_stderr = r_stderr
     self.tty = tty
+    self._stdout = sys.stdout
+    self._stderr = sys.stderr
   
   def __enter__(self):
+    # It's possible that __enter__ does not execute before __exit__ in some
+    # special cases. We also store _stdout and _stderr when creating the context.
     self._stdout = sys.stdout
     self._stderr = sys.stderr
     
     _remap_output_streams(self.r_stdout, self.r_stderr, self.tty)
   
-  def __exit__(self):
+  def __exit__(self, *args):
     sys.stdout = self._stdout
     sys.stderr = self._stderr
 
diff --git a/tests/testthat/_snaps/python-knitr-engine/knitr-print2.md b/tests/testthat/_snaps/python-knitr-engine/knitr-print2.md
new file mode 100644
index 000000000..330e35834
--- /dev/null
+++ b/tests/testthat/_snaps/python-knitr-engine/knitr-print2.md
@@ -0,0 +1,4 @@
+    bt <- reticulate::import_builtins()
+    bt$print("Hello world")
+
+    ## Hello world
diff --git a/tests/testthat/test-python-knitr-engine.R b/tests/testthat/test-python-knitr-engine.R
index d1d246e1a..a058e9012 100644
--- a/tests/testthat/test-python-knitr-engine.R
+++ b/tests/testthat/test-python-knitr-engine.R
@@ -71,7 +71,7 @@ test_that("Output streams are remaped when kniting", {
   local_edition(3)
 
   owd <- setwd(test_path("resources"))
-  rmarkdown::render("knitr-print.Rmd")
+  rmarkdown::render("knitr-print.Rmd", quiet = TRUE)
   setwd(owd)
 
   rendered <- test_path("resources", "knitr-print.md")
@@ -87,4 +87,15 @@ test_that("Output streams are remaped when kniting", {
   }))
   expect_length(out, 0)
 
+  # make sure it works from a background session
+  callr::r(function(path) {
+    setwd(path)
+    rmarkdown::render(
+      "knitr-print.Rmd",
+      output_file = "knitr-print2.md",
+      quiet = TRUE
+    )
+  }, args = list(path = test_path("resources")))
+
+  expect_snapshot_file(test_path("resources", "knitr-print2.md"))
 })