From 73e9250fd847dc4148b9d777a963616c85d26327 Mon Sep 17 00:00:00 2001
From: Daniel Wennberg <daniel.wennberg@gmail.com>
Date: Wed, 14 Aug 2024 08:27:06 -0700
Subject: [PATCH] Set IOContext for redirected stdout/stderr (#2824)

Co-authored-by: Reuben Gardos Reid <5456207+ReubenJ@users.noreply.github.com>
---
 src/runner/PlutoRunner/src/PlutoRunner.jl | 19 +++++++++----
 test/Logging.jl                           | 34 +++++++++++++++++++++++
 2 files changed, 47 insertions(+), 6 deletions(-)

diff --git a/src/runner/PlutoRunner/src/PlutoRunner.jl b/src/runner/PlutoRunner/src/PlutoRunner.jl
index dbf52ac55d..d39ec94aa3 100644
--- a/src/runner/PlutoRunner/src/PlutoRunner.jl
+++ b/src/runner/PlutoRunner/src/PlutoRunner.jl
@@ -1007,10 +1007,17 @@ const default_iocontext = IOContext(devnull,
     :pluto_with_js_link => (io, callback, on_cancellation) -> core_with_js_link(io, callback, on_cancellation),
 )
 
+# `stdout` mimics a TTY, the only relevant property is :color
 const default_stdout_iocontext = IOContext(devnull, 
-    :color => true, 
-    :limit => true, 
-    :displaysize => (18, 75), 
+    :color => true,
+    :is_pluto => false,
+)
+
+# `display` sees a richer context like in the REPL, see #2727
+const default_display_iocontext = IOContext(devnull,
+    :color => true,
+    :limit => true,
+    :displaysize => (18, 75),
     :is_pluto => false,
 )
 
@@ -2792,8 +2799,8 @@ function with_io_to_logs(f::Function; enabled::Bool=true, loglevel::Logging.LogL
     # Redirect both the `stdout` and `stderr` streams to a single `Pipe` object.
     pipe = Pipe()
     Base.link_pipe!(pipe; reader_supports_async = true, writer_supports_async = true)
-    pe_stdout = pipe.in
-    pe_stderr = pipe.in
+    pe_stdout = IOContext(pipe.in, default_stdout_iocontext)
+    pe_stderr = IOContext(pipe.in, default_stdout_iocontext)
     redirect_stdout(pe_stdout)
     redirect_stderr(pe_stderr)
 
@@ -2826,7 +2833,7 @@ function with_io_to_logs(f::Function; enabled::Bool=true, loglevel::Logging.LogL
     end
 
     # To make the `display` function work.
-    redirect_display = TextDisplay(IOContext(pe_stdout, default_stdout_iocontext))
+    redirect_display = TextDisplay(IOContext(pe_stdout, default_display_iocontext))
     pushdisplay(redirect_display)
 
     # Run the function `f`, capturing all output that it might have generated.
diff --git a/test/Logging.jl b/test/Logging.jl
index 8eb274bb54..c34d3f2654 100644
--- a/test/Logging.jl
+++ b/test/Logging.jl
@@ -59,6 +59,13 @@ using Pluto.WorkspaceManager: poll
         end
         """,
         "StructWithCustomShowThatLogs()", # 21
+        """
+        printstyled(stdout, "hello", color=:red)
+        """, # 22
+        "show(collect(1:500))", # 23
+        "show(stdout, collect(1:500))", # 24
+        "show(stdout, \"text/plain\", collect(1:500))", # 25
+        "display(collect(1:500))", # 26
     ]))
 
     @testset "Stdout" begin
@@ -123,6 +130,33 @@ using Pluto.WorkspaceManager: poll
         end
     end
 
+    @testset "ANSI Color Output" begin
+        update_run!(🍭, notebook, notebook.cells[22])
+        msg = only(notebook.cells[22].logs)["msg"][1]
+
+        @test startswith(msg, Base.text_colors[:red])
+        @test endswith(msg, Base.text_colors[:default])
+    end
+    
+    @testset "show(...) and display(...) behavior" begin
+        update_run!(🍭, notebook, notebook.cells[23:25])
+
+        msgs_show = [only(cell.logs)["msg"][1] for cell in notebook.cells[23:25]]
+
+        # `show` should show a middle element of the big array
+        for msg in msgs_show
+            @test contains(msg, "1") && contains(msg, "500")
+            @test contains(msg, "250")
+        end
+
+        update_run!(🍭, notebook, notebook.cells[26])
+        msg_display = only(notebook.cells[26].logs)["msg"][1]
+
+        # `display` should not display the middle element of the big array
+        @test contains(msg_display, "1") && contains(msg_display, "500")
+        @test !contains(msg_display, "250")
+    end
+
     @testset "Logging respects maxlog" begin
         @testset "Single log" begin