environment: Allow setting environment variables with -E.

* guix/scripts/environment.scm (show-environment-options-help) [-E]: Document
the change.
* guix/scripts/environment.scm (%options) [-E]: If arg contains =, split it
and store under 'extra-env.
* guix/scripts/environment.scm (launch-environment, launch-environment/fork)
(launch-environment/container): Take and use `environment-variables' keyword argument.
* guix/scripts/environment.scm (guix-environment*): Process 'environment-variables opts.
* guix/profiles.scm (load-profile): Add `environment-variables' keyword argument.
* doc/guix.texi (Invoking guix shell): Document the change.
* tests/guix-environment.sh (GUIX_TEST_ZZZ): Test it.

Change-Id: Ib1db19ef73dda8f561b05731079a96eda076e365
Signed-off-by: Maxim Cournoyer <maxim@guixotic.coop>
Modified-by: Maxim Cournoyer <maxim@guixotic.coop>

4 files changed, 63 insertions, 21 deletions

diff --git a/doc/guix.texi b/doc/guix.texi
index c624c497d2a..05eac9bda96 100644
--- a/doc/guix.texi
+++ b/doc/guix.texi
@@ -6803,20 +6803,26 @@ creating an environment in which search paths only contain package inputs.
 
 @item --preserve=@var{regexp}
 @itemx -E @var{regexp}
+@itemx -E @var{var}=@var{value}
 When used alongside @option{--pure}, preserve the environment variables
 matching @var{regexp}---in other words, put them on a ``white list'' of
-environment variables that must be preserved.  This option can be repeated
-several times.
+environment variables that must be preserved.  This option can be
+repeated several times.
+
+When @var{regexp} contains a variable assignment (with a @samp{=}
+character), the @var{var} environment variable is directly set to the
+@var{value} in the environment.  This is used in the example below to
+set the number of threads OpenMPI should use:
 
 @example
-guix shell --pure --preserve=^SLURM openmpi @dots{} \
+guix shell --pure --preserve=^SLURM -E OMP_NUM_THREADS=16 openmpi @dots{} \
   -- mpirun @dots{}
 @end example
 
 This example runs @command{mpirun} in a context where the only environment
 variables defined are @env{PATH}, environment variables whose name starts
-with @samp{SLURM}, as well as the usual ``precious'' variables (@env{HOME},
-@env{USER}, etc.).
+with @samp{SLURM}, variable @env{TERM} set to @samp{no=color}, as well as
+the usual ``precious'' variables (@env{HOME}, @env{USER}, etc.).
 
 @item --search-paths
 Display the environment variable definitions that make up the
@@ -7319,6 +7325,7 @@ creating an environment in which search paths only contain package inputs.
 
 @item --preserve=@var{regexp}
 @itemx -E @var{regexp}
+@itemx -E @var{var}=@var{value}
 When used alongside @option{--pure}, preserve the environment variables
 matching @var{regexp}---in other words, put them on a ``white list'' of
 environment variables that must be preserved.  This option can be repeated
diff --git a/guix/profiles.scm b/guix/profiles.scm
index b39116bc0bf..aced7bd25cb 100644
--- a/guix/profiles.scm
+++ b/guix/profiles.scm
@@ -2171,14 +2171,20 @@ WHITE-LIST-REGEXPS and those listed in WHITE-LIST."
 (define* (load-profile profile
                        #:optional (manifest (profile-manifest profile))
                        #:key pure? (white-list-regexps '())
+                       (environment-variables '())
                        (white-list %precious-variables))
   "Set the environment variables specified by MANIFEST for PROFILE.  When
 PURE? is #t, unset the variables in the current environment except those that
 match the regexps in WHITE-LIST-REGEXPS and those listed in WHITE-LIST.
 Otherwise, augment existing environment variables with additional search
-paths."
+paths.  Environment variables in ENVIRONMENT-VARIABLES, a list of pairs, are
+@emph{set} in the environment."
   (when pure?
     (purify-environment white-list white-list-regexps))
+  (map (match-lambda
+         ((name . value)
+          (setenv name value)))
+       environment-variables)
   (for-each (match-lambda
               ((($ <search-path-specification> variable _ separator) . value)
                (let ((current (getenv variable)))
diff --git a/guix/scripts/environment.scm b/guix/scripts/environment.scm
index dc7f3b9b280..bd7d4ca6ed4 100644
--- a/guix/scripts/environment.scm
+++ b/guix/scripts/environment.scm
@@ -100,7 +100,9 @@ shell'."
   (display (G_ "
       --pure             unset existing environment variables"))
   (display (G_ "
-  -E, --preserve=REGEXP  preserve environment variables that match REGEXP"))
+  -E, --preserve=REGEXP
+  -E VAR=VALUE           preserve environment variables that match REGEXP;
+                         if REGEXP contains `=', set VAR to VALUE instead"))
   (display (G_ "
       --search-paths     display needed environment variable definitions"))
   (display (G_ "
@@ -212,9 +214,15 @@ COMMAND or an interactive shell in that environment.\n"))
                    (alist-cons 'pure #t result)))
          (option '(#\E "preserve") #t #f
                  (lambda (opt name arg result)
-                   (alist-cons 'inherit-regexp
-                               (make-regexp* arg)
-                               result)))
+                   (let ((=-index (string-index arg #\=)))
+                     (if =-index
+                         (alist-cons 'environment-variables
+                                     (cons (substring arg 0 =-index)
+                                           (substring arg (1+ =-index)))
+                                     result)
+                         (alist-cons 'inherit-regexp
+                                     (make-regexp* arg)
+                                     result)))))
          (option '("inherit") #t #f               ;deprecated
                  (lambda (opt name arg result)
                    (warning (G_ "'--inherit' is deprecated, \
@@ -513,19 +521,23 @@ and suitable for 'exit'."
 
 (define* (launch-environment command profile manifest
                              #:key pure? (white-list '())
+                             (environment-variables '())
                              emulate-fhs?)
   "Load the environment of PROFILE, which corresponds to MANIFEST, and execute
 COMMAND.  When PURE?, pre-existing environment variables are cleared before
-setting the new ones, except those matching the regexps in WHITE-LIST.  When
-EMULATE-FHS?, first set up an FHS environment with $PATH and generate the LD
-cache."
+setting the new ones, except those matching the regexps in WHITE-LIST.
+Variables in ENVIRONMENT-VARIABLES (a list of pairs) are set in the
+environment.  When EMULATE-FHS?, first set up an FHS environment with $PATH
+and generate the LD cache."
   ;; Properly handle SIGINT, so pressing C-c in an interactive terminal
   ;; application works.
   (sigaction SIGINT SIG_DFL)
   ;; Restore original action for SIGPIPE.
   (sigaction SIGPIPE SIG_DFL)
   (load-profile profile manifest
-                #:pure? pure? #:white-list-regexps white-list)
+                #:pure? pure?
+                #:white-list-regexps white-list
+                #:environment-variables environment-variables)
 
   ;; Give users a way to know that they're in 'guix environment', so they can
   ;; adjust 'PS1' accordingly, for instance.  Set it to PROFILE so users can
@@ -734,15 +746,19 @@ command name."
                           closest))))))))
 
 (define* (launch-environment/fork command profile manifest
-                                  #:key pure? (white-list '()))
+                                  #:key
+                                  pure? (white-list '())
+                                  (environment-variables '()))
   "Run COMMAND in a new process with an environment containing PROFILE, with
 the search paths specified by MANIFEST.  When PURE?, pre-existing environment
 variables are cleared before setting the new ones, except those matching the
-regexps in WHITE-LIST."
+regexps in WHITE-LIST.  Variables in ENVIRONMENT-VARIABLES (a list of pairs)
+are set in the environment."
   (match (primitive-fork)
     (0 (launch-environment command profile manifest
                            #:pure? pure?
-                           #:white-list white-list))
+                           #:white-list white-list
+                           #:environment-variables environment-variables))
     (pid (match (waitpid pid)
            ((_ . status)
             status)))))
@@ -752,7 +768,8 @@ regexps in WHITE-LIST."
                                        map-cwd? cwd emulate-fhs? nesting?
                                        writable-root?
                                        (setup-hook #f)
-                                       (symlinks '()) (white-list '()))
+                                       (symlinks '()) (white-list '())
+                                       (environment-variables '()))
   "Run COMMAND within a container that features the software in PROFILE.
 Environment variables are set according to the search paths of MANIFEST.  The
 global shell is BASH, a file name for a GNU Bash binary in the store.  When
@@ -777,7 +794,10 @@ SYMLINKS must be a list of (SOURCE -> TARGET) tuples denoting symlinks to be
 added to the container.
 
 Preserve environment variables whose name matches the one of the regexps in
-WHILE-LIST."
+WHILE-LIST.
+
+Variables in ENVIRONMENT-VARIABLES (a list of pairs) are set in the
+environment."
   (define (optional-mapping->fs mapping)
     (and (file-exists? (file-system-mapping-source mapping))
          (file-system-mapping->bind-mount mapping)))
@@ -977,6 +997,7 @@ WHILE-LIST."
                                      (string-append home-dir "/.guix-profile")
                                      profile)
                                  manifest #:pure? #f
+                                 #:environment-variables environment-variables
                                  #:emulate-fhs? emulate-fhs?)))
           #:populate-file-system
           (lambda ()
@@ -1171,7 +1192,8 @@ command-line option processing with 'parse-command-line'."
                            '("/bin/sh")
                            (list %default-shell))))
          (mappings   (pick-all opts 'file-system-mapping))
-         (white-list (pick-all opts 'inherit-regexp)))
+         (white-list (pick-all opts 'inherit-regexp))
+         (environment-variables  (pick-all opts 'environment-variables)))
 
     (define store-needed?
       ;; Whether connecting to the daemon is needed.
@@ -1291,6 +1313,8 @@ when using '--container'; doing nothing~%"))
                                                   #:profile profile
                                                   #:manifest manifest
                                                   #:white-list white-list
+                                                  #:environment-variables
+                                                  environment-variables
                                                   #:link-profile? link-prof?
                                                   #:network? network?
                                                   #:map-cwd? (not no-cwd?)
@@ -1308,6 +1332,8 @@ when using '--container'; doing nothing~%"))
                    (exit/status
                     (launch-environment/fork command profile manifest
                                              #:white-list white-list
+                                             #:environment-variables
+                                             environment-variables
                                              #:pure? pure?)))))))))))))
 
 ;;; Local Variables:
diff --git a/tests/guix-environment.sh b/tests/guix-environment.sh
index 56cb289f2ff..450f0265221 100644
--- a/tests/guix-environment.sh
+++ b/tests/guix-environment.sh
@@ -49,18 +49,21 @@ test -x `sed -r 's/^export PATH="(.*)"/\1/' "$tmpdir/a"`/guile
 
 cmp "$tmpdir/a" "$tmpdir/b"
 
-# Check '--preserve'.
+# Check '--preserve', as well as its -E VAR=VALUE special case.
 GUIX_TEST_ABC=1
 GUIX_TEST_DEF=2
 GUIX_TEST_XYZ=3
+GUIX_TEST_ZZZ=4
 export GUIX_TEST_ABC GUIX_TEST_DEF GUIX_TEST_XYZ
 guix environment --bootstrap --ad-hoc guile-bootstrap --pure    \
      --preserve='^GUIX_TEST_A' --preserve='^GUIX_TEST_D'        \
+     --preserve='^GUIX_TEST_Z' -E 'GUIX_TEST_ZZZ=overridden'    \
      -- "$SHELL" -c set > "$tmpdir/a"
 grep '^PATH=' "$tmpdir/a"
 grep '^GUIX_TEST_ABC=1' "$tmpdir/a"
 grep '^GUIX_TEST_DEF=2' "$tmpdir/a"
 grep '^GUIX_TEST_XYZ=3' "$tmpdir/a" && false
+grep '^GUIX_TEST_ZZZ=overridden' "$tmpdir/a"
 
 # Make sure the exit value is preserved.
 if guix environment --bootstrap --ad-hoc guile-bootstrap --pure \