Add io_context keyword argument (#26)

Author: MichaelHatherly

CHANGELOG.md

@@ -3,6 +3,12 @@
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## Version [v0.2.5] - 2024-06-12
+
+### Added
+
+* `iocapture` now accepts an `io_context` keyword argument that is passed to the `IOContext` that wraps `stdout` and `stderr`. ([#26])
+
 ## Version [v0.2.4] - 2023-01-17
 
 ### Added
@@ -65,6 +71,7 @@ Initial release exporting the `iocapture` function.
 [v0.2.2]: https://github.com/JuliaDocs/IOCapture.jl/releases/tag/v0.2.2
 [v0.2.3]: https://github.com/JuliaDocs/IOCapture.jl/releases/tag/v0.2.3
 [v0.2.4]: https://github.com/JuliaDocs/IOCapture.jl/releases/tag/v0.2.4
+[v0.2.5]: https://github.com/JuliaDocs/IOCapture.jl/releases/tag/v0.2.5
 [#1]: https://github.com/JuliaDocs/IOCapture.jl/issues/1
 [#2]: https://github.com/JuliaDocs/IOCapture.jl/issues/2
 [#3]: https://github.com/JuliaDocs/IOCapture.jl/issues/3
@@ -76,4 +83,5 @@ Initial release exporting the `iocapture` function.
 [#20]: https://github.com/JuliaDocs/IOCapture.jl/issues/20
 [#21]: https://github.com/JuliaDocs/IOCapture.jl/issues/21
 [#23]: https://github.com/JuliaDocs/IOCapture.jl/issues/23
+[#26]: https://github.com/JuliaDocs/IOCapture.jl/issues/26
 [fredrikekre/Literate.jl#138]: https://github.com/fredrikekre/Literate.jl/issues/138

Project.toml

@@ -1,7 +1,7 @@
 name = "IOCapture"
 uuid = "b5f81e59-6552-4d32-b1f0-c071b021bf89"
 authors = ["Morten Piibeleht", "Michael Hatherly", "IOCapture contributors"]
-version = "0.2.4"
+version = "0.2.5"
 
 [deps]
 Logging = "56ddb016-857b-54e1-b83d-db4d58db5568"

src/IOCapture.jl

@@ -4,7 +4,7 @@ import Random
 
 """
     IOCapture.capture(
-        f; rethrow=Any, color=false, passthrough=false, capture_buffer=IOBuffer()
+        f; rethrow=Any, color=false, passthrough=false, capture_buffer=IOBuffer(), io_context=[],
     )
 
 Runs the function `f` and captures the `stdout` and `stderr` outputs, without printing
@@ -39,6 +39,10 @@ The behaviour can be customized with the following keyword arguments:
 * `capture_buffer`: The internal buffer used to capture the combined `stdout`
   and `stderr`.
 
+* `io_context`: An optional vector of `IOContext` key/value pairs that are passed to
+  the `IOContext` that wraps the redirected `stdout` and `stderr` streams. This only
+  has an effect on Julia v1.6 and later.
+
 # Extended help
 
 `capture` works by temporarily redirecting the standard output and error streams
@@ -98,8 +102,13 @@ function capture(
     rethrow::Type=Any,
     color::Bool=false,
     passthrough::Bool=false,
-    capture_buffer=IOBuffer()
+    capture_buffer=IOBuffer(),
+    io_context::AbstractVector=[],
 )
+    if any(x -> !isa(x, Pair{Symbol,<:Any}), io_context)
+        throw(ArgumentError("`io_context` must be a `Vector` of `Pair{Symbol,<:Any}`."))
+    end
+
     # Original implementation from Documenter.jl (MIT license)
     # Save the default output streams.
     default_stdout = stdout
@@ -109,8 +118,8 @@ function capture(
     pipe = Pipe()
     Base.link_pipe!(pipe; reader_supports_async = true, writer_supports_async = true)
     @static if VERSION >= v"1.6.0-DEV.481" # https://github.com/JuliaLang/julia/pull/36688
-        pe_stdout = IOContext(pipe.in, :color => get(stdout, :color, false) & color)
-        pe_stderr = IOContext(pipe.in, :color => get(stderr, :color, false) & color)
+        pe_stdout = IOContext(pipe.in, :color => get(stdout, :color, false) & color, io_context...)
+        pe_stderr = IOContext(pipe.in, :color => get(stderr, :color, false) & color, io_context...)
     else
         pe_stdout = pipe.in
         pe_stderr = pipe.in

test/runtests.jl

@@ -294,5 +294,25 @@ end
         end
     end
 
+    if VERSION >= v"1.6.0-DEV.481"
+        @testset "io_context" begin
+            # Avoids needing to define this at top-level as a `module M` syntax.
+            M = Module(:M)
+            Core.eval(M, :(struct T end))
 
+            no_io_context = IOCapture.capture() do
+                println(M.T())
+            end
+            @test rstrip(no_io_context.output) == "Main.M.T()"
+
+            with_io_context = IOCapture.capture(io_context=[:module => M]) do
+                println(M.T())
+            end
+            @test rstrip(with_io_context.output) == "T()"
+
+            @test_throws ArgumentError IOCapture.capture(io_context=["module" => M]) do
+                println(M.T())
+            end
+        end
+    end
 end