Implement `redirect_std(out|err)(::Function, ::TeeStream)`, closes #1.

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
 

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"