Allow key-deletions and atom keys in System.put_env/1

Author: josevalim

lib/elixir/lib/system.ex

@@ -733,11 +733,28 @@ defmodule System do
   Sets multiple environment variables.
 
   Sets a new value for each environment variable corresponding
-  to each `{key, value}` pair in `enum`.
+  to each `{key, value}` pair in `enum`. Keys are automatically
+  converted to strings, values are sent as is. `nil` values erase
+  the given keys.
   """
   @spec put_env(Enumerable.t()) :: :ok
   def put_env(enum) do
-    Enum.each(enum, fn {key, val} -> put_env(key, val) end)
+    Enum.each(enum, fn
+      {key, nil} ->
+        :os.unsetenv(to_charlist(key))
+
+      {key, val} ->
+        key = to_charlist(key)
+
+        case :string.find(key, "=") do
+          :nomatch ->
+            :os.putenv(key, to_charlist(val))
+
+          _ ->
+            raise ArgumentError,
+                  "cannot execute System.put_env/1 for key with \"=\", got: #{inspect(key)}"
+        end
+    end)
   end
 
   @doc """

lib/elixir/test/elixir/system_test.exs

@@ -50,7 +50,7 @@ defmodule SystemTest do
 
   @test_var "SYSTEM_ELIXIR_ENV_TEST_VAR"
 
-  test "*_env/*" do
+  test "get_env/put_env/delete_env" do
     assert System.get_env(@test_var) == nil
     assert System.get_env(@test_var, "SAMPLE") == "SAMPLE"
     assert System.fetch_env(@test_var) == :error
@@ -68,14 +68,22 @@ defmodule SystemTest do
     System.delete_env(@test_var)
     assert System.get_env(@test_var) == nil
 
-    System.put_env(%{@test_var => "OTHER_SAMPLE"})
-    assert System.get_env(@test_var) == "OTHER_SAMPLE"
-
     assert_raise ArgumentError, ~r[cannot execute System.put_env/2 for key with \"=\"], fn ->
       System.put_env("FOO=BAR", "BAZ")
     end
   end
 
+  test "put_env/2" do
+    System.put_env(%{@test_var => "MAP_STRING"})
+    assert System.get_env(@test_var) == "MAP_STRING"
+
+    System.put_env([{String.to_atom(@test_var), "KW_ATOM"}])
+    assert System.get_env(@test_var) == "KW_ATOM"
+
+    System.put_env([{String.to_atom(@test_var), nil}])
+    assert System.get_env(@test_var) == nil
+  end
+
   test "cmd/2 raises for null bytes" do
     assert_raise ArgumentError, ~r"cannot execute System.cmd/3 for program with null byte", fn ->
       System.cmd("null\0byte", [])

lib/mix/lib/mix.ex

@@ -567,9 +567,8 @@ defmodule Mix do
       apps
 
     * `:system_env` (since v1.13.0) - a list or a map of system environment variable
-      names as binary keys and their respective values as binaries. The system environment
-      is made part of the `Mix.install/2` cache, so different configurations will lead
-      to different apps
+      names with respective values as binaries. The system environment is made part
+      of the `Mix.install/2` cache, so different configurations will lead to different apps
 
   ## Examples
 
@@ -589,7 +588,7 @@ defmodule Mix do
           nx: [default_backend: EXLA]
         ],
         system_env: [
-          {"XLA_TARGET", "cuda111"}
+          XLA_TARGET: "cuda111"
         ]
       )
 

lib/mix/lib/mix/dep/loader.ex

@@ -61,20 +61,15 @@ defmodule Mix.Dep.Loader do
   def with_system_env(%Mix.Dep{system_env: new_env}, callback) do
     old_env =
       for {key, _} <- new_env do
+        key = to_string(key)
         {key, System.get_env(key)}
       end
 
     try do
       System.put_env(new_env)
       callback.()
     after
-      for {key, value} <- old_env do
-        if value do
-          System.put_env(key, value)
-        else
-          System.delete_env(key)
-        end
-      end
+      System.put_env(old_env)
     end
   end
 
@@ -208,7 +203,7 @@ defmodule Mix.Dep.Loader do
       requirement: req,
       status: scm_status(scm, opts),
       opts: Keyword.put_new(opts, :env, :prod),
-      system_env: system_env
+      system_env: Enum.to_list(system_env)
     }
   end
 

lib/mix/test/mix_test.exs

@@ -141,13 +141,15 @@ defmodule MixTest do
           {:install_test, path: Path.join(tmp_dir, "install_test")}
         ],
         config: [unknown_app: [foo: :bar]],
-        system_env: %{"MIX_INSTALL_FOO" => "BAR"}
+        system_env: %{"MIX_INSTALL_FOO" => "BAR", MIX_INSTALL_BAZ: "BAT"}
       )
 
       assert Application.fetch_env!(:unknown_app, :foo) == :bar
       assert System.fetch_env!("MIX_INSTALL_FOO") == "BAR"
+      assert System.fetch_env!("MIX_INSTALL_BAZ") == "BAT"
     after
       System.delete_env("MIX_INSTALL_FOO")
+      System.delete_env("MIX_INSTALL_BAZ")
       Application.delete_env(:unknown_app, :foo, persistent: true)
     end