From 438e02529dfb4b26603e2a06992df977cd80ecd9 Mon Sep 17 00:00:00 2001
From: Samuel Dionne-Riel <samuel@dionne-riel.com>
Date: Wed, 1 Aug 2018 21:28:49 -0400
Subject: [PATCH] Allows selectively adding environment variables to pure
 shells.

Includes documentation and test.
---
 doc/manual/command-ref/nix-shell.xml | 8 ++++++++
 src/nix-build/nix-build.cc           | 7 ++++++-
 tests/nix-shell.sh                   | 7 +++++++
 3 files changed, 21 insertions(+), 1 deletion(-)

diff --git a/doc/manual/command-ref/nix-shell.xml b/doc/manual/command-ref/nix-shell.xml
index 62d026ac2..5c44c4a8f 100644
--- a/doc/manual/command-ref/nix-shell.xml
+++ b/doc/manual/command-ref/nix-shell.xml
@@ -32,6 +32,7 @@
     <arg><option>--run</option> <replaceable>cmd</replaceable></arg>
     <arg><option>--exclude</option> <replaceable>regexp</replaceable></arg>
     <arg><option>--pure</option></arg>
+    <arg><option>--keep</option> <replaceable>name</replaceable></arg>
     <group choice='req'>
       <arg choice='plain'>
         <group choice='req'>
@@ -165,6 +166,13 @@ also <xref linkend="sec-common-options" />.</phrase></para>
 
     </listitem></varlistentry>
 
+  <varlistentry><term><option>--keep</option> <replaceable>name</replaceable></term>
+
+    <listitem><para>When a <option>--pure</option> shell is started,
+    keep the listed environment variables.</para></listitem>
+
+  </varlistentry>
+
 </variablelist>
 
 <para>The following common options are supported:</para>
diff --git a/src/nix-build/nix-build.cc b/src/nix-build/nix-build.cc
index de0e9118f..a7ec206aa 100755
--- a/src/nix-build/nix-build.cc
+++ b/src/nix-build/nix-build.cc
@@ -99,6 +99,9 @@ void mainWrapped(int argc, char * * argv)
 
     std::string outLink = "./result";
 
+    // List of environment variables kept for --pure
+    std::set<string> keepVars{"HOME", "USER", "LOGNAME", "DISPLAY", "PATH", "TERM", "IN_NIX_SHELL", "TZ", "PAGER", "NIX_BUILD_SHELL", "SHLVL"};
+
     Strings args;
     for (int i = 1; i < argc; ++i)
         args.push_back(argv[i]);
@@ -218,6 +221,9 @@ void mainWrapped(int argc, char * * argv)
             }
         }
 
+        else if (*arg == "--keep")
+            keepVars.insert(getArg(*arg, arg, end));
+
         else if (*arg == "-")
             readStdin = true;
 
@@ -368,7 +374,6 @@ void mainWrapped(int argc, char * * argv)
         auto tmp = getEnv("TMPDIR", getEnv("XDG_RUNTIME_DIR", "/tmp"));
 
         if (pure) {
-            std::set<string> keepVars{"HOME", "USER", "LOGNAME", "DISPLAY", "PATH", "TERM", "IN_NIX_SHELL", "TZ", "PAGER", "NIX_BUILD_SHELL", "SHLVL"};
             decltype(env) newEnv;
             for (auto & i : env)
                 if (keepVars.count(i.first))
diff --git a/tests/nix-shell.sh b/tests/nix-shell.sh
index d25c456ce..6024ea399 100644
--- a/tests/nix-shell.sh
+++ b/tests/nix-shell.sh
@@ -4,12 +4,19 @@ clearStore
 
 # Test nix-shell -A
 export IMPURE_VAR=foo
+export SELECTED_IMPURE_VAR=baz
 export NIX_BUILD_SHELL=$SHELL
 output=$(nix-shell --pure shell.nix -A shellDrv --run \
     'echo "$IMPURE_VAR - $VAR_FROM_STDENV_SETUP - $VAR_FROM_NIX"')
 
 [ "$output" = " - foo - bar" ]
 
+# Test --keep
+output=$(nix-shell --pure --keep SELECTED_IMPURE_VAR shell.nix -A shellDrv --run \
+    'echo "$IMPURE_VAR - $VAR_FROM_STDENV_SETUP - $VAR_FROM_NIX - $SELECTED_IMPURE_VAR"')
+
+[ "$output" = " - foo - bar - baz" ]
+
 # Test nix-shell on a .drv
 [[ $(nix-shell --pure $(nix-instantiate shell.nix -A shellDrv) --run \
     'echo "$IMPURE_VAR - $VAR_FROM_STDENV_SETUP - $VAR_FROM_NIX"') = " - foo - bar" ]]