more ergonomic stream redirection

NEWS.md

@@ -44,6 +44,7 @@ New library functions
 * Two arguments method `lock(f, lck)` now accepts a `Channel` as the second argument. ([#39312])
 * New functor `Returns(value)`, which returns `value` for any arguments ([#39794])
 * New macro `Base.@invoke f(arg1::T1, arg2::T2; kwargs...)` provides an easier syntax to call `invoke(f, Tuple{T1,T2}, arg1, arg2; kwargs...)` ([#38438])
+* New function `redirect` for redirecting `stdin`, `stdout` and `stderr` ([#37978]).
 
 New library features
 --------------------

base/exports.jl

@@ -827,6 +827,7 @@ export
     readline,
     readlines,
     readuntil,
+    redirect,
     redirect_stderr,
     redirect_stdin,
     redirect_stdout,

base/stream.jl

@@ -1204,6 +1204,8 @@ the pipe.
 !!! note
     `stream` must be a compatible objects, such as an `IOStream`, `TTY`,
     `Pipe`, socket, or `devnull`.
+
+See also [`redirect`](@ref).
 """
 redirect_stdout
 
@@ -1215,6 +1217,8 @@ Like [`redirect_stdout`](@ref), but for [`stderr`](@ref).
 !!! note
     `stream` must be a compatible objects, such as an `IOStream`, `TTY`,
     `Pipe`, socket, or `devnull`.
+
+See also [`redirect`](@ref).
 """
 redirect_stderr
 
@@ -1227,9 +1231,118 @@ Note that the direction of the stream is reversed.
 !!! note
     `stream` must be a compatible objects, such as an `IOStream`, `TTY`,
     `Pipe`, socket, or `devnull`.
+
+See also [`redirect`](@ref).
 """
 redirect_stdin
 
+"""
+    redirect(;stdin=stdin, stderr=stderr, stdout=stdout)
+
+Redirect a subset of the streams `stdin`, `stderr`, `stdout`.
+Each argument must be an `IOStream`, `TTY`, `Pipe`, socket, or `devnull`.
+"""
+function redirect(;stdin=nothing, stderr=nothing, stdout=nothing)
+    stdin  === nothing || redirect_stdin(stdin)
+    stderr === nothing || redirect_stderr(stderr)
+    stdout === nothing || redirect_stdout(stdout)
+end
+
+"""
+    redirect(f; stdin=nothing, stderr=nothing, stdout=nothing)
+
+Redirect a subset of the streams `stdin`, `stderr`, `stdout`,
+call `f()` and restore each stream.
+
+Possible values for each stream are:
+* `nothing` indicating the stream should not be redirected.
+* `path::AbstractString` redirecting the stream to the file at `path`.
+* `io` an `IOStream`, `TTY`, `Pipe`, socket, or `devnull`.
+
+# Examples
+```julia
+julia> redirect(stdout="stdout.txt", stderr="stderr.txt") do
+           print("hello stdout")
+           print(stderr, "hello stderr")
+       end
+
+julia> read("stdout.txt", String)
+"hello stdout"
+
+julia> read("stderr.txt", String)
+"hello stderr"
+```
+
+# Edge cases
+
+It is possible to pass the same argument to `stdout` and `stderr`:
+```julia
+julia> redirect(stdout="log.txt", stderr="log.txt", stdin=devnull) do
+    ...
+end
+```
+
+However it is not supported to pass two distinct descriptors of the same file.
+```julia
+julia> io1 = open("same/path", "w")
+
+julia> io2 = open("same/path", "w")
+
+julia> redirect(f, stdout=io1, stderr=io2) # not suppored
+```
+Also the `stdin` argument may not be the same descriptor as `stdout` or `stderr`.
+```julia
+julia> io = open(...)
+
+julia> redirect(f, stdout=io, stdin=io) # not supported
+```
+
+!!! compat "Julia 1.6"
+    `redirect` requires Julia 1.6 or later.
+"""
+function redirect(f; stdin=nothing, stderr=nothing, stdout=nothing)
+
+    function resolve(new::Nothing, oldstream, mode)
+        (new=nothing, close=false, old=nothing)
+    end
+    function resolve(path::AbstractString, oldstream,mode)
+        (new=open(path, mode), close=true, old=oldstream)
+    end
+    function resolve(new, oldstream, mode)
+        (new=new, close=false, old=oldstream)
+    end
+
+    same_path(x, y) = false
+    same_path(x::AbstractString, y::AbstractString) = x == y
+    if same_path(stderr, stdin)
+        throw(ArgumentError("stdin and stderr cannot be the same path"))
+    end
+    if same_path(stdout, stdin)
+        throw(ArgumentError("stdin and stdout cannot be the same path"))
+    end
+
+    new_in , close_in , old_in  = resolve(stdin , Base.stdin , "r")
+    new_out, close_out, old_out = resolve(stdout, Base.stdout, "w")
+    if same_path(stderr, stdout)
+        # make sure that in case stderr = stdout = "same/path"
+        # only a single io is used instead of opening the same file twice
+        new_err, close_err, old_err = new_out, false, Base.stderr
+    else
+        new_err, close_err, old_err = resolve(stderr, Base.stderr, "w")
+    end
+
+    redirect(; stderr=new_err, stdin=new_in, stdout=new_out)
+
+    try
+        return f()
+    finally
+        redirect(;stderr=old_err, stdin=old_in, stdout=old_out)
+        close_err && close(new_err)
+        close_in  && close(new_in )
+        close_out && close(new_out)
+    end
+end
+
 function (f::redirect_stdio)(thunk::Function, stream)
     stdold = f.unix_fd == 0 ? stdin :
              f.unix_fd == 1 ? stdout :
@@ -1243,6 +1356,7 @@ function (f::redirect_stdio)(thunk::Function, stream)
     end
 end
 
+
 """
     redirect_stdout(f::Function, stream)
 

doc/src/base/io-network.md

@@ -36,6 +36,7 @@ Base.iswritable
 Base.isreadable
 Base.isopen
 Base.fd
+Base.redirect
 Base.redirect_stdout
 Base.redirect_stdout(::Function, ::Any)
 Base.redirect_stderr

test/spawn.jl

@@ -261,6 +261,70 @@ end
     end
 end
 
+@testset "redirect" begin
+
+    function hello_err_out()
+        println(stderr, "hello from stderr")
+        println(stdout, "hello from stdout")
+    end
+    @testset "same path for multiple streams" begin
+        @test_throws ArgumentError redirect(hello_err_out,
+                                            stdin="samepath.txt", stdout="samepath.txt")
+        @test_throws ArgumentError redirect(hello_err_out,
+                                            stdin="samepath.txt", stderr="samepath.txt")
+        mktempdir() do dir
+            path = joinpath(dir, "stdouterr.txt")
+            redirect(hello_err_out, stdout=path, stderr=path)
+            @test read(path, String) == """
+            hello from stderr
+            hello from stdout
+            """
+        end
+    end
+
+    mktempdir() do dir
+        path_stdout = joinpath(dir, "stdout.txt")
+        path_stderr = joinpath(dir, "stderr.txt")
+        redirect(hello_err_out, stderr=devnull, stdout=path_stdout)
+        @test read(path_stdout, String) == "hello from stdout\n"
+
+        open(path_stderr, "w") do ioerr
+            redirect(hello_err_out, stderr=ioerr, stdout=devnull)
+        end
+        @test read(path_stderr, String) == "hello from stderr\n"
+    end
+
+    mktempdir() do dir
+        path_stderr = joinpath(dir, "stderr.txt")
+        path_stdin  = joinpath(dir, "stdin.txt")
+        path_stdout = joinpath(dir, "stdout.txt")
+
+        content_stderr = randstring()
+        content_stdout = randstring()
+
+        redirect(stdout=path_stdout, stderr=path_stderr) do
+            print(content_stdout)
+            print(stderr, content_stderr)
+        end
+
+        @test read(path_stderr, String) == content_stderr
+        @test read(path_stdout, String) == content_stdout
+    end
+
+    # stdin is unavailable on the workers. Run test on master.
+    ret = Core.eval(Main,
+            quote
+                remotecall_fetch(1) do
+                    mktempdir() do dir
+                        path = joinpath(dir, "stdin.txt")
+                        write(path, "hello from stdin\n")
+                        redirect(readline, stdin=path)
+                    end
+                end
+            end)
+    @test ret == "hello from stdin"
+end
+
 # issue #36136
 @testset "redirect to devnull" begin
     @test redirect_stdout(devnull) do; println("Hello") end === nothing