feat: add WithEnv method for setting a task's environment variables

script.go

@@ -34,6 +34,12 @@ type Pipe struct {
 	// because pipe stages are concurrent, protect 'err'
 	mu  *sync.Mutex
 	err error
+
+	// env contains the environment to run any exec commands with.
+	// Each entry in the array should be of the form key=value.
+	// If env is not nil, it will replace the default environment variables
+	// when executing commands.
+	env []string
 }
 
 // Args creates a pipe containing the program's command-line arguments from
@@ -166,6 +172,7 @@ func NewPipe() *Pipe {
 		mu:         new(sync.Mutex),
 		stdout:     os.Stdout,
 		httpClient: http.DefaultClient,
+		env:        nil,
 	}
 }
 
@@ -388,6 +395,9 @@ func (p *Pipe) Exec(cmdLine string) *Pipe {
 		if p.stderr != nil {
 			cmd.Stderr = p.stderr
 		}
+		if p.env != nil {
+			cmd.Env = p.env
+		}
 		err = cmd.Start()
 		if err != nil {
 			fmt.Fprintln(cmd.Stderr, err)
@@ -428,6 +438,9 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe {
 			if p.stderr != nil {
 				cmd.Stderr = p.stderr
 			}
+			if p.env != nil {
+				cmd.Env = p.env
+			}
 			err = cmd.Start()
 			if err != nil {
 				fmt.Fprintln(cmd.Stderr, err)
@@ -898,6 +911,15 @@ func (p *Pipe) WithStdout(w io.Writer) *Pipe {
 	return p
 }
 
+// WithEnv sets the environment for the exec commands to the string array env.
+// This will override the default process environment variables when executing
+// commands run via [Pipe.Exec] or [Pipe.ExecForEach]. This will not affect the
+// environment outside of [Pipe.Exec] or [Pipe.ExecForEach].
+func (p *Pipe) WithEnv(env []string) *Pipe {
+	p.env = env
+	return p
+}
+
 // WriteFile writes the pipe's contents to the file path, truncating it if it
 // exists, and returns the number of bytes successfully written, or an error.
 func (p *Pipe) WriteFile(path string) (int64, error) {

script_test.go

@@ -1768,6 +1768,52 @@ func TestWithStdout_SetsSpecifiedWriterAsStdout(t *testing.T) {
 	}
 }
 
+func TestWithEnv_UnsetsAllEnvVarsGivenEmptySlice(t *testing.T) {
+	t.Parallel()
+	want := []string{}
+
+	output, err := script.NewPipe().WithEnv(want).Exec("printenv").String()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if output != "" {
+		t.Errorf("want empty environment, got %q", output)
+	}
+}
+
+func TestWithEnv_SetMultipleEnvVars(t *testing.T) {
+	t.Parallel()
+	env := []string{"ENV1=test1", "ENV2=test2"}
+
+	got, err := script.NewPipe().WithEnv(env).Exec("printenv").String()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	want := "ENV1=test1\nENV2=test2\n"
+
+	if got != want {
+		t.Errorf("want %v, got %v", want, got)
+	}
+}
+
+func TestNotSettingEnvFallsBackToDefaultEnvironment(t *testing.T) {
+	t.Parallel()
+
+	got, err := script.NewPipe().Exec("printenv").String()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// not setting the environment should simply use the task's default environment
+	want := strings.Join(os.Environ(), "\n") + "\n"
+
+	if got != want {
+		t.Errorf("want %v, got %v", want, got)
+	}
+}
+
 func TestErrorReturnsErrorSetByPreviousPipeStage(t *testing.T) {
 	t.Parallel()
 	p := script.File("testdata/nonexistent.txt")