diff --git a/src/TeeStreams.jl b/src/TeeStreams.jl index 5983c29..b450aa9 100644 --- a/src/TeeStreams.jl +++ b/src/TeeStreams.jl @@ -120,6 +120,37 @@ end # end # end +# Adapted from Logging2.jl (MIT License: Copyright (c) 2020 Chris Foster) +# https://github.com/JuliaLogging/Logging2.jl/blob/094eb6619aeaa8815585dbe7f33b4972f9a4ce6b/src/Logging2.jl#L12-L37 +function (redirect_func::Base.RedirectStdStream)(f::Function, io::TeeStream) + prev_stream = + redirect_func.unix_fd == 1 ? stdout : + redirect_func.unix_fd == 2 ? stderr : + throw(ArgumentError("Can only redirect stdout and stderr to TeeStream.")) + + result = nothing + + rd, rw = redirect_func() + try + @sync begin + try + Threads.@spawn write(io, rd) # loops while !eof(rd) + result = f() + finally + # To close the read side of the pipe, we must close *all* + # writers. This includes `rw`, but *also* the dup'd fd + # created behind the scenes by redirect_func(). (To close + # that, must call redirect_func() here with the prev stream.) + close(rw) + redirect_func(prev_stream) + end + end + finally + close(rd) + end + return result +end + end # module diff --git a/test/runtests.jl b/test/runtests.jl index e163570..a4430f0 100644 --- a/test/runtests.jl +++ b/test/runtests.jl @@ -67,6 +67,22 @@ end @test read(f1, String) == read(f2, String) == correct end + # Redirection of std(out|err) to TeeStream + io = IOBuffer() + ioc = IOContext(IOBuffer()) + tee = TeeStream(io, ioc) + ret1, ret2 = redirect_stderr(tee) do + r1 = redirect_stdout(tee) do + print(stderr, "stderr") + print(stdout, "stdout") + return 1 + end + return r1, 2 + end + @test ret1 == 1 + @test ret2 == 2 + @test String(take!(io)) == String(take!(ioc.io)) == "stderrstdout" + # Test some integration with other packages mktempdir() do tmpd url = "https://julialang-s3.julialang.org/bin/linux/x64/1.5/julia-1.5.3-linux-x86_64.tar.gz"