script: Allow setting env var value from stdout

script/cmds.go

@@ -449,11 +449,17 @@ func Echo() Cmd {
 // With no arguments, Env reports all variables in the environment.
 // "key=value" arguments set variables, and arguments without "="
 // cause the corresponding value to be printed to the stdout buffer.
+//
+// Passing --from-stdout will set the "key" to the stdout of the previous
+// command as the "value".
 func Env() Cmd {
 	return Command(
 		CmdUsage{
 			Summary: "set or log the values of environment variables",
-			Args:    "[key[=value]...]",
+			Flags: func(fs *pflag.FlagSet) {
+				fs.Bool("from-stdout", false, "Whether to set the value of the variable from stdout")
+			},
+			Args: "[key[=value]...]",
 			Detail: []string{
 				"With no arguments, print the script environment to the log.",
 				"Otherwise, add the listed key=value pairs to the environment or print the listed keys.",
@@ -474,21 +480,32 @@ func Env() Cmd {
 			},
 		},
 		func(s *State, args ...string) (WaitFunc, error) {
+			fromStdout, err := s.Flags.GetBool("from-stdout")
+			if err != nil {
+				return nil, err
+			}
+
 			out := new(strings.Builder)
 			if len(args) == 0 {
 				for _, kv := range s.env {
 					fmt.Fprintf(out, "%s\n", kv)
 				}
 			} else {
 				for _, env := range args {
-					i := strings.Index(env, "=")
-					if i < 0 {
-						// Display value instead of setting it.
-						fmt.Fprintf(out, "%s=%s\n", env, s.envMap[env])
-						continue
-					}
-					if err := s.Setenv(env[:i], env[i+1:]); err != nil {
-						return nil, err
+					if fromStdout {
+						if err := s.Setenv(env, strings.Trim(s.Stdout(), "\n ")); err != nil {
+							return nil, err
+						}
+					} else {
+						i := strings.Index(env, "=")
+						if i < 0 {
+							// Display value instead of setting it.
+							fmt.Fprintf(out, "%s=%s\n", env, s.envMap[env])
+							continue
+						}
+						if err := s.Setenv(env[:i], env[i+1:]); err != nil {
+							return nil, err
+						}
 					}
 				}
 			}

script/scripttest/scripttest.go

@@ -103,19 +103,7 @@ func Run(t testing.TB, e *script.Engine, s *script.State, filename string, testS
 
 		if testing.Verbose() {
 			// Add the environment to the start of the script log.
-			wait, err := script.Env().Run(s)
-			if err != nil {
-				t.Fatal(err)
-			}
-			if wait != nil {
-				stdout, stderr, err := wait(s)
-				if err != nil {
-					t.Fatalf("env: %v\n%s", err, stderr)
-				}
-				if len(stdout) > 0 {
-					s.Logf("%s\n", stdout)
-				}
-			}
+			err = e.ExecuteLine(s, "env", log)
 		}
 
 		return e.Execute(s, filename, bufio.NewReader(testScript), log)

script/scripttest/testdata/env.txt

@@ -0,0 +1,29 @@
+# Tests for the env command
+
+# set an env var
+env TEST=value
+env TEST
+cmp stdout stdout_env.txt
+
+# set an env var from stdout
+echo hello
+env --from-stdout TEST
+env TEST
+cmp stdout stdout_env_stdout.txt
+
+# set multiple env vars from stdout
+echo world
+env --from-stdout TEST1 TEST2
+env TEST1
+cmp stdout stdout_env_stdout_1.txt
+env TEST2
+cmp stdout stdout_env_stdout_2.txt
+
+-- stdout_env.txt --
+TEST=value
+-- stdout_env_stdout.txt --
+TEST=hello
+-- stdout_env_stdout_1.txt --
+TEST1=world
+-- stdout_env_stdout_2.txt --
+TEST2=world