Merge pull request #10914 from stevengj/with_env

export Base.withenv
JuliaLang · Apr 20, 2015 · a49ae53 · vtjnash · Apr 21, 2015 · tkelman
2 parents f65befe + 468f640
commit a49ae53
Show file tree

Hide file tree

Showing 8 changed files with 38 additions and 15 deletions.
diff --git a/NEWS.md b/NEWS.md
@@ -271,6 +271,9 @@ Library improvements
     * `code_llvm` now outputs stripped IR without debug info or other attached metadata.
       Use `code_llvm_raw` for the unstripped output ([#10747]).
 
+    * New `withenv(var=>val, ...) do ... end` function to temporarily
+      modify environment variables ([#10914]).
+
 Deprecated or removed
 ---------------------
 
@@ -1371,3 +1374,5 @@ Too numerous to mention.
 [#10747]: https://github.com/JuliaLang/julia/issues/10747
 [#10844]: https://github.com/JuliaLang/julia/issues/10844
 [#10870]: https://github.com/JuliaLang/julia/issues/10870
+[#10885]: https://github.com/JuliaLang/julia/issues/10885
+[#10914]: https://github.com/JuliaLang/julia/issues/10914
diff --git a/base/deprecated.jl b/base/deprecated.jl
@@ -442,6 +442,8 @@ end
 @deprecate cholfact(A::AbstractMatrix, β::Number) cholfact(A, shift=β)
 @deprecate ldltfact(A::AbstractMatrix, β::Number) ldltfact(A, shift=β)
 
+@deprecate with_env(f::Function, key::AbstractString, val) withenv(f, key=>val)
+
 # 0.4 discontinued functions
 
 @noinline function subtypetree(x::DataType, level=-1)

diff --git a/base/env.jl b/base/env.jl
@@ -169,16 +169,20 @@ function show(io::IO, ::EnvHash)
 end
 
 # temporarily set and then restore an environment value
-function with_env(f::Function, key::AbstractString, val)
-    old = get(ENV,key,nothing)
-    val != nothing ? (ENV[key]=val) : _unsetenv(key)
+function withenv{T<:AbstractString}(f::Function, keyvals::Pair{T}...)
+    old = Dict{T,Any}()
+    for (key,val) in keyvals
+        old[key] = get(ENV,key,nothing)
+        val != nothing ? (ENV[key]=val) : _unsetenv(key)
+    end
     try f()
     finally
-        old != nothing ? (ENV[key]=old) : _unsetenv(key)
-    catch
-        rethrow()
+        for (key,val) in old
+            val != nothing ? (ENV[key]=val) : _unsetenv(key)
+        end
     end
 end
+withenv(f::Function) = f() # handle empty keyvals case; see #10853
 
 ## misc environment-related functionality ##
 

diff --git a/base/exports.jl b/base/exports.jl
@@ -1299,6 +1299,7 @@ export
     setenv,
     spawn,
     success,
+    withenv,
 
 # C interface
     cfunction,

diff --git a/base/pkg/entry.jl b/base/pkg/entry.jl
@@ -246,7 +246,7 @@ function update(branch::AbstractString)
             Git.run(`checkout -q $branch`)
         end
         # TODO: handle merge conflicts
-        Base.with_env("GIT_MERGE_AUTOEDIT","no") do
+        Base.withenv("GIT_MERGE_AUTOEDIT"=>"no") do
             Git.run(`pull --rebase -q`, out=DevNull)
         end
     end

diff --git a/base/process.jl b/base/process.jl
@@ -119,6 +119,7 @@ detach(cmd::Cmd) = (cmd.detach=true; cmd)
 
 setenv{S<:ByteString}(cmd::Cmd, env::Array{S}; dir="") = (cmd.env = ByteString[x for x in env]; setenv(cmd, dir=dir); cmd)
 setenv(cmd::Cmd, env::Associative; dir="") = (cmd.env = ByteString[string(k)*"="*string(v) for (k,v) in env]; setenv(cmd, dir=dir); cmd)
+setenv{T<:AbstractString}(cmd::Cmd, env::Pair{T}...; dir="") = (cmd.env = ByteString[k*"="*string(v) for (k,v) in env]; setenv(cmd, dir=dir); cmd)
 setenv(cmd::Cmd; dir="") = (cmd.dir = dir; cmd)
 
 (&)(left::AbstractCmd, right::AbstractCmd) = AndCmds(left, right)

diff --git a/doc/stdlib/base.rst b/doc/stdlib/base.rst
@@ -676,12 +676,19 @@ System
 
 .. function:: setenv(command, env; dir=working_dir)
 
-   Set environment variables to use when running the given command. ``env`` is either
-   a dictionary mapping strings to strings, or an array of strings of the form
-   ``"var=val"``.
+   Set environment variables to use when running the given
+   command. ``env`` is either a dictionary mapping strings to strings,
+   an array of strings of the form ``"var=val"``, or zero or more
+   ``"var"=>val`` pair arguments.  In order to modify (rather than
+   replace) the existing environment, create ``env`` by ``copy(ENV)``
+   and then setting ``env["var"]=val`` as desired, or use ``withenv``.
 
-   The ``dir`` keyword argument can be used to specify a working directory for the
-   command.
+   The ``dir`` keyword argument can be used to specify a working
+   directory for the command.
+
+.. function:: withenv(f::Function, kv::Pair...)
+
+   Execute ``f()`` in an environment that is temporarily modified (not replaced as in ``setenv``) by zero or more ``"var"=>val`` arguments ``kv``.  ``withenv`` is generally used via the ``withenv(kv...) do ... end`` syntax.  A value of ``nothing`` can be used to temporarily unset an environment variable (if it is set).  When ``withenv`` returns, the original environment has been restored.
 
 .. function:: pipe(from, to, ...)
 

diff --git a/test/spawn.jl b/test/spawn.jl
@@ -92,9 +92,12 @@ rm(file)
     end
 end
 
-readall(setenv(`sh -c "echo \$TEST"`,["TEST=Hello World"])) == "Hello World\n"
-readall(setenv(`sh -c "echo \$TEST"`,Dict("TEST"=>"Hello World"))) == "Hello World\n"
-readall(setenv(`sh -c "pwd"`;dir="/")) == readall(setenv(`sh -c "cd / && pwd"`))
+@test readall(setenv(`sh -c "echo \$TEST"`,["TEST=Hello World"])) == "Hello World\n"
+@test readall(setenv(`sh -c "echo \$TEST"`,Dict("TEST"=>"Hello World"))) == "Hello World\n"
+@test readall(setenv(`sh -c "echo \$TEST"`,"TEST"=>"Hello World")) == "Hello World\n"
+@test (withenv("TEST"=>"Hello World") do
+       readall(`sh -c "echo \$TEST"`); end) == "Hello World\n"
+@test readall(setenv(`sh -c "pwd"`;dir="..")) == readall(setenv(`sh -c "cd .. && pwd"`))
 
 # Here we test that if we close a stream with pending writes, we don't lose the writes.
 str = ""