Add color keyword argument

README.md

@@ -21,6 +21,8 @@ See the docstring for full documentation.
 
 ## Known limitations
 
+### Separately stored `stdout` or `stderr` objects
+
 The capturing does not work properly if `f` prints to the `stdout` object that has been
 stored in a separate variable or object, e.g.:
 
@@ -56,6 +58,13 @@ Stacktrace:
 This is because `stdout` and `stderr` within an `iocapture` actually refer to the temporary
 redirect streams which get cleaned up at the end of the `iocapture` call.
 
+### ANSI color/escape code
+
+On Julia 1.5 and earlier, setting `color` to `true` has no effect, because the [ability to set `IOContext` attributes on
+redirected streams was added in 1.6](https://github.com/JuliaLang/julia/pull/36688). I.e. on those older Julia versions
+the captured output will generally not contain ANSI color escape sequences.
+
+
 ## Similar packages
 
 * [Suppressor.jl](https://github.com/JuliaIO/Suppressor.jl) provides similar functionality,

src/IOCapture.jl

@@ -4,7 +4,7 @@ using Logging
 export iocapture
 
 """
-    iocapture(f; throwerrors=true)
+    iocapture(f; throwerrors=true, color=false)
 
 Runs the function `f` and captures the `stdout` and `stderr` outputs without printing them
 in the terminal. Returns an object with the following fields:
@@ -21,6 +21,10 @@ The behaviour can be customized with the following keyword arguments:
   via the `.value` field (with also `.error` and `.backtrace` set accordingly). If set to
   `:interrupt`, only `InterruptException`s are rethrown.
 
+* `color`: if set to `true`, `iocapture` inherits the `:color` property of `stdout` and
+  `stderr`, which specifies whether ANSI color/escape codes are expected. This argument is
+  only effective on Julia v1.6 and later.
+
 # Extended help
 
 `iocapture` works by temporarily redirecting the standard output and error streams
@@ -49,7 +53,7 @@ It is also possible to set `throwerrors = :interrupt`, which will make `iocaptur
 only `InterruptException`s. This is useful when you want to capture all the exceptions, but
 allow the user to interrupt the running code with `Ctrl+C`.
 """
-function iocapture(f; throwerrors::Union{Bool,Symbol}=true)
+function iocapture(f; throwerrors::Union{Bool,Symbol}=true, color::Bool=false)
     # Currently, :interrupt is the only valid Symbol value for throwerrors
     if isa(throwerrors, Symbol) && throwerrors !== :interrupt
         throw(DomainError(throwerrors, "Invalid value passed for throwerrors"))
@@ -62,10 +66,17 @@ function iocapture(f; throwerrors::Union{Bool,Symbol}=true)
     # 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)
-    redirect_stdout(pipe.in)
-    redirect_stderr(pipe.in)
+    @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)
+    else
+        pe_stdout = pipe.in
+        pe_stderr = pipe.in
+    end
+    redirect_stdout(pe_stdout)
+    redirect_stderr(pe_stderr)
     # Also redirect logging stream to the same pipe
-    logger = ConsoleLogger(pipe.in)
+    logger = ConsoleLogger(pe_stderr)
 
     # Bytes written to the `pipe` are captured in `output` and converted to a `String`.
     output = UInt8[]

test/runtests.jl

@@ -5,6 +5,10 @@ using Test
 # Note: this can not be inside the testset
 (VERSION < v"1.2.0-DEV.272") && (hasfield(::Type{T}, name::Symbol) where T = Base.fieldindex(T, name, false) > 0)
 
+hascolor(io) = VERSION >= v"1.6.0-DEV.481" && get(io, :color, false)
+has_escapecodes(s) = occursin(r"\e\[[^m]*m", s)
+strip_escapecodes(s) = replace(s, r"\e\[[^m]*m" => "")
+
 # Callable object for testing
 struct Foo
     x
@@ -68,8 +72,20 @@ end
     @test c.output == "foo"
     @test c.value === nothing
 
+    # Colors are preserved if it's supported
+    c = iocapture(color=true) do
+        printstyled("foo", color=:red)
+    end
+    @test !c.error
+    if hascolor(stdout)
+        @test c.output == "\e[31mfoo\e[39m"
+    else
+        @test c.output == "foo"
+    end
+    @test c.value === nothing
+
     # This test checks that deprecation warnings are captured correctly
-    c = iocapture() do
+    c = iocapture(color=true) do
         println("println")
         @info "@info"
         f() = (Base.depwarn("depwarn", :f); nothing)
@@ -83,9 +99,12 @@ end
     @test isdefined(Base, :JLOptions)
     @test hasfield(Base.JLOptions, :depwarn)
     if Base.JLOptions().depwarn == 0 # --depwarn=no, default on Julia >= 1.5
-        @test c.output == "println\n[ Info: @info\n"
+        @test has_escapecodes(c.output) === hascolor(stderr)
+        @test strip_escapecodes(c.output) == "println\n[ Info: @info\n"
     else # --depwarn=yes
-        @test startswith(c.output, "println\n[ Info: @info\n┌ Warning: depwarn\n")
+        @test has_escapecodes(c.output) === hascolor(stderr)
+        output_nocol = strip_escapecodes(c.output)
+        @test startswith(output_nocol, "println\n[ Info: @info\n┌ Warning: depwarn\n")
     end
 
     # Exceptions -- normally rethrown