Allow pass-through of output (#20)

CHANGELOG.md

@@ -1,5 +1,9 @@
 # IOCapture.jl changelog
 
+## Unreleased
+
+* ![Enhancement][badge-enhancement] `iocapture` now accepts a `passthrough` keyword argument that passes through output to `stdout` as well as capturing it. ([#19][github-19], [#20][github-20])
+
 ## Version `0.2.3`
 
 * ![Bugfix][badge-bugfix] User code that creates a lot of "method definition overwritten" warnings no longer stalls in `IOCapture.capture` due to a buffer not being emptied. ([JuliaDocs/Documenter.jl#2121][documenter-2121], [#15][github-15])
@@ -44,6 +48,8 @@ Initial release exporting the `iocapture` function.
 [github-9]: https://github.com/JuliaDocs/IOCapture.jl/pull/9
 [github-11]: https://github.com/JuliaDocs/IOCapture.jl/pull/11
 [github-15]: https://github.com/JuliaDocs/IOCapture.jl/pull/15
+[github-19]: https://github.com/JuliaDocs/IOCapture.jl/issues/19
+[github-20]: https://github.com/JuliaDocs/IOCapture.jl/pull/20
 
 [literate-138]: https://github.com/fredrikekre/Literate.jl/issues/138
 

src/IOCapture.jl

@@ -3,10 +3,12 @@ using Logging
 import Random
 
 """
-    IOCapture.capture(f; rethrow=Any, color=false)
+    IOCapture.capture(f; rethrow=Any, color=false, passthrough=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:
+Runs the function `f` and captures the `stdout` and `stderr` outputs, without printing
+them in the terminal, unless `passthrough=true`.
+
+Returns an object with the following fields:
 
 * `.value :: Any`: return value of the function, or the error exception object on error
 * `.output :: String`: captured `stdout` and `stderr`
@@ -48,6 +50,11 @@ julia> c.output
 
 This approach does have some limitations -- see the README for more information.
 
+If `passthrough=true`, the redirected streams will also be passed through to the
+original standard output. As a result, the output from `f` would both be captured and
+shown on screen. Note that `stdout` and `stderr` are merged in the pass-through, and color
+is stripped unless the `color` option is set to `true`.
+
 **Exceptions.** Normally, if `f` throws an exception, `capture` simply re-throws it with
 `rethrow`. However, by setting `rethrow`, it is also possible to capture errors, which then
 get returned via the `.value` field. Additionally, `.error` is set to `true`, to indicate
@@ -69,7 +76,7 @@ using IOcapture: capture as iocapture
 
 This avoids the function name being too generic.
 """
-function capture(f; rethrow::Type=Any, color::Bool=false)
+function capture(f; rethrow::Type=Any, color::Bool=false, passthrough::Bool=false)
     # Original implementation from Documenter.jl (MIT license)
     # Save the default output streams.
     default_stdout = stdout
@@ -105,7 +112,20 @@ function capture(f; rethrow::Type=Any, color::Bool=false)
     # pipe to `output` in order to avoid the buffer filling up and stalling write() calls in
     # user code.
     output = IOBuffer()
-    buffer_redirect_task = @async write(output, pipe)
+    if passthrough
+        bufsize = 128
+        buffer = Vector{UInt8}(undef, bufsize)
+        buffer_redirect_task = @async begin
+            while !eof(pipe)
+                nbytes = readbytes!(pipe, buffer, bufsize)
+                data = view(buffer, 1:nbytes)
+                write(output, data)
+                write(default_stdout, data)
+            end
+        end
+    else
+        buffer_redirect_task = @async write(output, pipe)
+    end
 
     if old_rng !== nothing
         copy!(Random.default_rng(), old_rng)

test/runtests.jl

@@ -212,4 +212,49 @@ end
         end
         @test true # just make sure we get here
     end
+
+    @testset "passthrough" begin
+        mktemp() do logfile, io
+            redirect_stdout(io) do
+                print("<pre>")
+                c = IOCapture.capture(passthrough=true) do
+                    for i in 1:128
+                        print("HelloWorld")
+                    end
+                end
+                print("<post>")
+            end
+            close(io)
+            @test c.output == "HelloWorld"^128
+            @test read(logfile, String) == "<pre>" * "HelloWorld"^128 * "<post>"
+        end
+        # Interaction of passthrough= with color=
+        # Also tests that stdout and stderr get merged in both .output and passthrough
+        if VERSION >= v"1.6.0"
+            # older versions don't support `redirect_stdout(IOContext…`
+            mktemp() do logfile, io
+                redirect_stdout(IOContext(io, :color => true)) do
+                    c = IOCapture.capture(passthrough=true) do
+                        printstyled(stdout, "foo"; color=:blue)
+                        printstyled(stderr, "bar"; color=:red)
+                    end
+                end
+                close(io)
+                @test c.output == "foobar"
+                @test c.output == read(logfile, String)
+            end
+            mktemp() do logfile, io
+                redirect_stdout(IOContext(io, :color => true)) do
+                    c = IOCapture.capture(passthrough=true, color=true) do
+                        printstyled(stdout, "foo"; color=:blue)
+                        printstyled(stderr, "bar"; color=:red)
+                    end
+                end
+                close(io)
+                @test c.output == "\e[34mfoo\e[39m\e[31mbar\e[39m"
+                @test c.output == read(logfile, String)
+            end
+        end
+    end
+
 end