From 9152ab9ef07e4ae1b226097d25821a57a41703b6 Mon Sep 17 00:00:00 2001
From: Fredrik Ekre <ekrefredrik@gmail.com>
Date: Wed, 7 Jul 2021 11:45:53 +0200
Subject: [PATCH] Remove trace of sandbox module in output for at-repl and
 at-example.

---
 CHANGELOG.md               |  3 ++
 src/Expanders.jl           | 23 ++++++++++---
 test/examples/src/index.md | 66 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 88 insertions(+), 4 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 7d7ae8c307..2b6b6f5a68 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -4,6 +4,8 @@
 
 * ![Feature][badge-feature] `@example`- and `@repl`-blocks now support colored output by mapping ANSI escape sequences to HTML. This requires Julia >= 1.6 and passing `ansicolor=true` to `Documenter.HTML` (e.g. `makedocs(format=Documenter.HTML(ansicolor=true, ...), ...)`). In Documenter 0.28.0 this will be the default so to (preemptively) opt-out pass `ansicolor=false`. ([#1441][github-1441], [#1628][github-1628])
 
+* ![Enhancement][badge-enhancement] The sandbox module used for evaluating `@repl` and `@example` blocks is now removed (replaced with `Main`) in text output. ([#1633][github-1633])
+
 * ![Bugfix][badge-bugfix] Dollar signs in the HTML output no longer get accidentally misinterpreted as math delimiters in the browser. ([#890](github-890), [#1625](github-1625))
 
 ## Version `v0.27.3`
@@ -853,6 +855,7 @@
 [github-1617]: https://github.com/JuliaDocs/Documenter.jl/pull/1617
 [github-1625]: https://github.com/JuliaDocs/Documenter.jl/pull/1625
 [github-1628]: https://github.com/JuliaDocs/Documenter.jl/pull/1628
+[github-1633]: https://github.com/JuliaDocs/Documenter.jl/pull/1633
 
 [julia-38079]: https://github.com/JuliaLang/julia/issues/38079
 [julia-39841]: https://github.com/JuliaLang/julia/pull/39841
diff --git a/src/Expanders.jl b/src/Expanders.jl
index 2fdcedebfb..96bea1d219 100644
--- a/src/Expanders.jl
+++ b/src/Expanders.jl
@@ -597,12 +597,18 @@ function Selectors.runner(::Type{ExampleBlocks}, x, page, doc)
 
     # Generate different  in different formats and let each writer select
     output = Base.invokelatest(Utilities.display_dict, result, context = :color => ansicolor)
+    # Remove references to gensym'd module from text/plain
+    m = MIME"text/plain"()
+    if haskey(output, m)
+        output[m] = remove_sandbox_from_output(output[m], mod)
+    end
 
     # Only add content when there's actually something to add.
     isempty(input) || push!(content, Markdown.Code("julia", input))
     if result === nothing
-        code = Documenter.DocTests.sanitise(buffer)
-        isempty(code) || push!(content, Dict{MIME,Any}(MIME"text/plain"() => code))
+        stdouterr = Documenter.DocTests.sanitise(buffer)
+        stdouterr = remove_sandbox_from_output(stdouterr, mod)
+        isempty(stdouterr) || push!(content, Dict{MIME,Any}(MIME"text/plain"() => stdouterr))
     elseif !isempty(output)
         push!(content, output)
     end
@@ -610,6 +616,11 @@ function Selectors.runner(::Type{ExampleBlocks}, x, page, doc)
     page.mapping[x] = Documents.MultiOutput(content)
 end
 
+# Replace references to gensym'd module with Main
+function remove_sandbox_from_output(str, mod::Module)
+    replace(str, Regex(("(Main\\.)?$(nameof(mod))")) => "Main")
+end
+
 # @repl
 # -----
 
@@ -655,13 +666,17 @@ function Selectors.runner(::Type{REPLBlocks}, x, page, doc)
             push!(multicodeblock, Markdown.Code("julia-repl", prepend_prompt(input)))
         end
         out = IOBuffer()
-        print(out, c.output)
+        print(out, c.output) # c.output is std(out|err)
         if isempty(input) || isempty(output)
             println(out)
         else
             println(out, output, "\n")
         end
-        push!(multicodeblock, Markdown.Code("documenter-ansi", rstrip(String(take!(out)))))
+
+        outstr = String(take!(out))
+        # Replace references to gensym'd module with Main
+        outstr = remove_sandbox_from_output(outstr, mod)
+        push!(multicodeblock, Markdown.Code("documenter-ansi", rstrip(outstr)))
     end
     page.mapping[x] = Documents.MultiCodeBlock("julia-repl", multicodeblock)
 end
diff --git a/test/examples/src/index.md b/test/examples/src/index.md
index b49d014d28..55ae11f436 100644
--- a/test/examples/src/index.md
+++ b/test/examples/src/index.md
@@ -506,3 +506,69 @@ julia> function foo end;
 
 I will pay \$1 if $x^2$ is displayed correctly. People may also write \$s or
 even money bag\$\$.
+
+# Module scrubbing from `@repl` and `@example`
+
+None of these expressions should result in the gensym'd module in the output
+
+```@repl
+@__MODULE__
+println("@__MODULE__ is ", @__MODULE__) # sandbox printed to stdout
+function f()
+    println("@__MODULE__ is ", @__MODULE__)
+    @warn "Main as the module for this log message"
+    @__MODULE__
+end
+f()
+@warn "Main as the module for this log message"
+```
+```@repl
+module A
+    function f()
+        println("@__MODULE__ is ", @__MODULE__)
+        @warn "Main.A as the module for this log message"
+        @__MODULE__
+    end
+end
+A.f()
+```
+
+```@example
+@__MODULE__ # sandbox as return value
+```
+
+```@example
+println("@__MODULE__ is ", @__MODULE__) # sandbox printed to stdout
+```
+
+```@example
+function f()
+    println("@__MODULE__ is ", @__MODULE__)
+end
+f()
+```
+
+```@example
+function f()
+    @__MODULE__
+end
+f()
+```
+
+```@example
+@warn "Main as the module for this log message"
+```
+
+```@example moduleA
+module A
+    function f()
+        println("@__MODULE__ is ", @__MODULE__)
+        @warn "Main.A as the module for this log message"
+        @__MODULE__
+    end
+end
+```
+
+```@example moduleA
+A.f()
+```