more ergonomic stream redirection (JuliaLang#37978)

Co-authored-by: Jeff Bezanson <[email protected]> Co-authored-by: Fredrik Ekre <[email protected]> Co-authored-by: Kristoffer Carlsson <[email protected]>
shirodkara · Jun 9, 2021 · 99c7a56 · 99c7a56
1 parent 3164581
commit 99c7a56
Show file tree

Hide file tree

Showing 5 changed files with 200 additions and 9 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -57,6 +57,7 @@ New library functions
 * 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 macros `@something` and `@coalesce` which are short-circuiting versions of `something` and `coalesce`, respectively ([#40729])
+* New function `redirect_stdio` for redirecting `stdin`, `stdout` and `stderr` ([#37978]).
 
 New library features
 --------------------

diff --git a/base/exports.jl b/base/exports.jl
@@ -833,6 +833,7 @@ export
     readline,
     readlines,
     readuntil,
+    redirect_stdio,
     redirect_stderr,
     redirect_stdin,
     redirect_stdout,

diff --git a/base/stream.jl b/base/stream.jl
@@ -1127,15 +1127,15 @@ function _fd(x::Union{LibuvStream, LibuvServer})
     return fd[]
 end
 
-struct redirect_stdio <: Function
+struct RedirectStdStream <: Function
     unix_fd::Int
     writable::Bool
 end
 for (f, writable, unix_fd) in
         ((:redirect_stdin, false, 0),
          (:redirect_stdout, true, 1),
          (:redirect_stderr, true, 2))
-    @eval const ($f) = redirect_stdio($unix_fd, $writable)
+    @eval const ($f) = RedirectStdStream($unix_fd, $writable)
 end
 function _redirect_io_libc(stream, unix_fd::Int)
     posix_fd = _fd(stream)
@@ -1154,7 +1154,7 @@ function _redirect_io_global(io, unix_fd::Int)
     unix_fd == 2 && (global stderr = io)
     nothing
 end
-function (f::redirect_stdio)(handle::Union{LibuvStream, IOStream})
+function (f::RedirectStdStream)(handle::Union{LibuvStream, IOStream})
     _redirect_io_libc(handle, f.unix_fd)
     c_sym = f.unix_fd == 0 ? cglobal(:jl_uv_stdin, Ptr{Cvoid}) :
             f.unix_fd == 1 ? cglobal(:jl_uv_stdout, Ptr{Cvoid}) :
@@ -1164,31 +1164,31 @@ function (f::redirect_stdio)(handle::Union{LibuvStream, IOStream})
     _redirect_io_global(handle, f.unix_fd)
     return handle
 end
-function (f::redirect_stdio)(::DevNull)
+function (f::RedirectStdStream)(::DevNull)
     nulldev = @static Sys.iswindows() ? "NUL" : "/dev/null"
     handle = open(nulldev, write=f.writable)
     _redirect_io_libc(handle, f.unix_fd)
     close(handle) # handle has been dup'ed in _redirect_io_libc
     _redirect_io_global(devnull, f.unix_fd)
     return devnull
 end
-function (f::redirect_stdio)(io::AbstractPipe)
+function (f::RedirectStdStream)(io::AbstractPipe)
     io2 = (f.writable ? pipe_writer : pipe_reader)(io)
     f(io2)
     _redirect_io_global(io, f.unix_fd)
     return io
 end
-function (f::redirect_stdio)(p::Pipe)
+function (f::RedirectStdStream)(p::Pipe)
     if p.in.status == StatusInit && p.out.status == StatusInit
         link_pipe!(p)
     end
     io2 = getfield(p, f.writable ? :in : :out)
     f(io2)
     return p
 end
-(f::redirect_stdio)() = f(Pipe())
+(f::RedirectStdStream)() = f(Pipe())
 
-# Deprecate these in v2 (redirect_stdio support)
+# Deprecate these in v2 (RedirectStdStream support)
 iterate(p::Pipe) = (p.out, 1)
 iterate(p::Pipe, i::Int) = i == 1 ? (p.in, 2) : nothing
 getindex(p::Pipe, key::Int) = key == 1 ? p.out : key == 2 ? p.in : throw(KeyError(key))
@@ -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_stdio`](@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_stdio`](@ref).
 """
 redirect_stderr
 
@@ -1227,10 +1231,125 @@ 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_stdio`](@ref).
 """
 redirect_stdin
 
-function (f::redirect_stdio)(thunk::Function, stream)
+"""
+    redirect_stdio(;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`.
+
+!!! compat "Julia 1.7"
+    `redirect_stdio` requires Julia 1.7 or later.
+"""
+function redirect_stdio(;stdin=nothing, stderr=nothing, stdout=nothing)
+    stdin  === nothing || redirect_stdin(stdin)
+    stderr === nothing || redirect_stderr(stderr)
+    stdout === nothing || redirect_stdout(stdout)
+end
+
+"""
+    redirect_stdio(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_stdio(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_stdio(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_stdio(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_stdio(f, stdout=io, stdin=io) # not supported
+```
+
+!!! compat "Julia 1.7"
+    `redirect_stdio` requires Julia 1.7 or later.
+"""
+function redirect_stdio(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
+    function same_path(x::AbstractString, y::AbstractString)
+        # if x = y = "does_not_yet_exist.txt" then samefile will return false
+        (abspath(x) == abspath(y)) || samefile(x,y)
+    end
+    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_stdio(; stderr=new_err, stdin=new_in, stdout=new_out)
+
+    try
+        return f()
+    finally
+        redirect_stdio(;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::RedirectStdStream)(thunk::Function, stream)
     stdold = f.unix_fd == 0 ? stdin :
              f.unix_fd == 1 ? stdout :
              f.unix_fd == 2 ? stderr :
@@ -1243,6 +1362,7 @@ function (f::redirect_stdio)(thunk::Function, stream)
     end
 end
 
+
 """
     redirect_stdout(f::Function, stream)
 

diff --git a/doc/src/base/io-network.md b/doc/src/base/io-network.md
@@ -36,6 +36,7 @@ Base.iswritable
 Base.isreadable
 Base.isopen
 Base.fd
+Base.redirect_stdio
 Base.redirect_stdout
 Base.redirect_stdout(::Function, ::Any)
 Base.redirect_stderr

diff --git a/test/spawn.jl b/test/spawn.jl
@@ -261,6 +261,74 @@ end
     end
 end
 
+@testset "redirect_stdio" 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_stdio(hello_err_out,
+                                            stdin="samepath.txt", stdout="samepath.txt")
+        @test_throws ArgumentError redirect_stdio(hello_err_out,
+                                            stdin="samepath.txt", stderr="samepath.txt")
+
+        @test_throws ArgumentError redirect_stdio(hello_err_out,
+                                            stdin=joinpath("tricky", "..", "samepath.txt"),
+                                            stderr="samepath.txt")
+        mktempdir() do dir
+            path = joinpath(dir, "stdouterr.txt")
+            redirect_stdio(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_stdio(hello_err_out, stderr=devnull, stdout=path_stdout)
+        @test read(path_stdout, String) == "hello from stdout\n"
+
+        open(path_stderr, "w") do ioerr
+            redirect_stdio(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_stdio(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_stdio(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