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,14 @@ func (p *Pipe) WithStdout(w io.Writer) *Pipe {
 	return p
 }
 
+// WithEnv sets the pipe's environment to the string array env. This will override
+// the default process environment variables when executing commands run via [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

@@ -12,6 +12,7 @@ import (
 	"os"
 	"path/filepath"
 	"regexp"
+	"slices"
 	"strings"
 	"testing"
 	"testing/iotest"
@@ -1768,6 +1769,57 @@ func TestWithStdout_SetsSpecifiedWriterAsStdout(t *testing.T) {
 	}
 }
 
+func TestWithEnv_SetEmptyEnvironment(t *testing.T) {
+	t.Parallel()
+	buf := new(bytes.Buffer)
+	env := []string{}
+
+	_, err := script.NewPipe().WithStdout(buf).WithEnv(env).Exec("printenv").Stdout()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	if len(buf.String()) != 0 {
+		got := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
+		t.Errorf("expected 0 environment variables, got %d %v", len(got), got)
+	}
+}
+
+func TestWithEnv_SetMultipleEnvVars(t *testing.T) {
+	t.Parallel()
+	buf := new(bytes.Buffer)
+	env := []string{"ENV1=test1", "ENV2=test2"}
+
+	_, err := script.NewPipe().WithStdout(buf).WithEnv(env).Exec("printenv").Stdout()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	got := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
+
+	if !slices.Equal(got, env) {
+		t.Errorf("expected %v, got %v", env, got)
+	}
+}
+
+func TestNotSettingEnvFallsBackToDefaultEnvironment(t *testing.T) {
+	t.Parallel()
+	buf := new(bytes.Buffer)
+
+	_, err := script.NewPipe().WithStdout(buf).Exec("printenv").Stdout()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	// not setting the environment should simply use the task's default environment
+	expected := os.Environ()
+	got := strings.Split(strings.TrimRight(buf.String(), "\n"), "\n")
+
+	if !slices.Equal(got, expected) {
+		t.Errorf("expected %v, got %v", expected, got)
+	}
+}
+
 func TestErrorReturnsErrorSetByPreviousPipeStage(t *testing.T) {
 	t.Parallel()
 	p := script.File("testdata/nonexistent.txt")