From afb0239caa4a477b5a4304338457c393efbb6026 Mon Sep 17 00:00:00 2001
From: Grant Bruer <grant.git.email@gmail.com>
Date: Wed, 18 Sep 2024 15:28:44 -0400
Subject: [PATCH] Write display() calls in markdown output

---
 src/Literate.jl  | 71 ++++++++++++++++++++++++++++--------------------
 test/runtests.jl | 22 +++++++++++++++
 2 files changed, 64 insertions(+), 29 deletions(-)

diff --git a/src/Literate.jl b/src/Literate.jl
index 9ce052e..15dec7a 100644
--- a/src/Literate.jl
+++ b/src/Literate.jl
@@ -641,39 +641,21 @@ function execute_markdown!(io::IO, sb::Module, block::String, outputdir;
                            flavor::AbstractFlavor, image_formats::Vector, file_prefix::String,
                            softscope::Bool)
     # TODO: Deal with explicit display(...) calls
-    r, str, _ = execute_block(sb, block; inputfile=inputfile, fake_source=fake_source, softscope=softscope)
+    r, str, display_dicts = execute_block(sb, block; inputfile=inputfile, fake_source=fake_source, softscope=softscope)
+
     # issue #101: consecutive codefenced blocks need newline
     # issue #144: quadruple backticks allow for triple backticks in the output
     plain_fence = "\n````\n" =>  "\n````"
-    if r !== nothing && !REPL.ends_with_semicolon(block)
-        if (flavor isa FranklinFlavor || flavor isa DocumenterFlavor) &&
-           Base.invokelatest(showable, MIME("text/html"), r)
-            htmlfence = flavor isa FranklinFlavor ? ("~~~" => "~~~") : ("```@raw html" => "```")
-            write(io, "\n", htmlfence.first, "\n")
-            Base.invokelatest(show, io, MIME("text/html"), r)
-            write(io, "\n", htmlfence.second, "\n")
-            return
-        end
-        for (mime, ext) in image_formats
-            if Base.invokelatest(showable, mime, r)
-                file = file_prefix * ext
-                open(joinpath(outputdir, file), "w") do io
-                    Base.invokelatest(show, io, mime, r)
-                end
-                write(io, "![](", file, ")\n")
-                return
-            end
-        end
-        if Base.invokelatest(showable, MIME("text/markdown"), r)
-            write(io, '\n')
-            Base.invokelatest(show, io, MIME("text/markdown"), r)
-            write(io, '\n')
-            return
+
+    # Any explicit calls to display(...)
+    for dict in display_dicts
+        for (mime, data) in dict
+            display_markdown(io, data, outputdir, flavor, image_formats, file_prefix, plain_fence)
         end
-        # fallback to text/plain
-        write(io, plain_fence.first)
-        Base.invokelatest(show, io, "text/plain", r)
-        write(io, plain_fence.second, '\n')
+    end
+
+    if r !== nothing && !REPL.ends_with_semicolon(block)
+        display_markdown(io, r, outputdir, flavor, image_formats, file_prefix, plain_fence)
         return
     elseif !isempty(str)
         write(io, plain_fence.first, str, plain_fence.second, '\n')
@@ -681,6 +663,37 @@ function execute_markdown!(io::IO, sb::Module, block::String, outputdir;
     end
 end
 
+function display_markdown(io, data, outputdir, flavor, image_formats, file_prefix, plain_fence)
+    if (flavor isa FranklinFlavor || flavor isa DocumenterFlavor) &&
+        Base.invokelatest(showable, MIME("text/html"), data)
+        htmlfence = flavor isa FranklinFlavor ? ("~~~" => "~~~") : ("```@raw html" => "```")
+        write(io, "\n", htmlfence.first, "\n")
+        Base.invokelatest(show, io, MIME("text/html"), data)
+        write(io, "\n", htmlfence.second, "\n")
+        return
+    end
+    for (mime, ext) in image_formats
+        if Base.invokelatest(showable, mime, data)
+            file = file_prefix * ext
+            open(joinpath(outputdir, file), "w") do io
+                Base.invokelatest(show, io, mime, data)
+            end
+            write(io, "![](", file, ")\n")
+            return
+        end
+    end
+    if Base.invokelatest(showable, MIME("text/markdown"), data)
+        write(io, '\n')
+        Base.invokelatest(show, io, MIME("text/markdown"), data)
+        write(io, '\n')
+        return
+    end
+    # fallback to text/plain
+    write(io, plain_fence.first)
+    Base.invokelatest(show, io, "text/plain", data)
+    write(io, plain_fence.second, '\n')
+    return
+end
 
 const JUPYTER_VERSION = v"4.3.0"
 
diff --git a/test/runtests.jl b/test/runtests.jl
index 529231d..78b0224 100644
--- a/test/runtests.jl
+++ b/test/runtests.jl
@@ -985,6 +985,28 @@ end end
                               flavor=Literate.CommonMarkFlavor())
             @test read(joinpath(outdir, "inputfile-1.svg"), String) == "issue228"
 
+            # Calls to display(x) and display(mime, x)
+            script = """
+            struct DF x end
+            Base.show(io::IO, ::MIME"text/plain", df::DF) = print(io, "DF(\$(df.x)) as text/plain")
+            Base.show(io::IO, ::MIME"text/html", df::DF) = print(io, "DF(\$(df.x)) as text/html")
+            Base.show(io::IO, ::MIME"text/latex", df::DF) = print(io, "DF(\$(df.x)) as text/latex")
+            #-
+            foreach(display, [DF(1), DF(2)])
+            DF(3)
+            #-
+            display(MIME("text/latex"), DF(4))
+            """
+            write(inputfile, script)
+            Literate.markdown(inputfile, outdir; execute=true)
+            markdown = read(joinpath(outdir, "inputfile.md"), String)
+            @test occursin("````\n\"DF(1) as text/plain\"\n````", markdown)
+            @test occursin("````\n\"DF(1) as text/html\"\n````", markdown)
+            @test occursin("````\n\"DF(2) as text/plain\"\n````", markdown)
+            @test occursin("````\n\"DF(2) as text/html\"\n````", markdown)
+            @test occursin("```@raw html\nDF(3) as text/html\n```", markdown)
+            @test occursin("````\n\"DF(4) as text/latex\"\n````", markdown)
+
             # Softscope
             write(
                 inputfile,