add WithEnv for setting command environment (#208)

Co-authored-by: John Arundel <[email protected]>

Author: mahadzaryab1

README.md

@@ -269,18 +269,31 @@ These are functions that create a pipe with a given contents:
 
 | Source | Contents |
 | -------- | ------------- |
-| [`Args`](https://pkg.go.dev/github.com/bitfield/script#Args) | command-line arguments
-| [`Do`](https://pkg.go.dev/github.com/bitfield/script#Do) | HTTP response
-| [`Echo`](https://pkg.go.dev/github.com/bitfield/script#Echo) | a string
-| [`Exec`](https://pkg.go.dev/github.com/bitfield/script#Exec) | command output
-| [`File`](https://pkg.go.dev/github.com/bitfield/script#File) | file contents
-| [`FindFiles`](https://pkg.go.dev/github.com/bitfield/script#FindFiles) | recursive file listing
-| [`Get`](https://pkg.go.dev/github.com/bitfield/script#Get) | HTTP response
-| [`IfExists`](https://pkg.go.dev/github.com/bitfield/script#IfExists) | do something only if some file exists
-| [`ListFiles`](https://pkg.go.dev/github.com/bitfield/script#ListFiles) | file listing (including wildcards)
-| [`Post`](https://pkg.go.dev/github.com/bitfield/script#Post) | HTTP response
-| [`Slice`](https://pkg.go.dev/github.com/bitfield/script#Slice) | slice elements, one per line
-| [`Stdin`](https://pkg.go.dev/github.com/bitfield/script#Stdin) | standard input
+| [`Args`](https://pkg.go.dev/github.com/bitfield/script#Args) | command-line arguments |
+| [`Do`](https://pkg.go.dev/github.com/bitfield/script#Do) | HTTP response |
+| [`Echo`](https://pkg.go.dev/github.com/bitfield/script#Echo) | a string |
+| [`Exec`](https://pkg.go.dev/github.com/bitfield/script#Exec) | command output |
+| [`File`](https://pkg.go.dev/github.com/bitfield/script#File) | file contents |
+| [`FindFiles`](https://pkg.go.dev/github.com/bitfield/script#FindFiles) | recursive file listing |
+| [`Get`](https://pkg.go.dev/github.com/bitfield/script#Get) | HTTP response |
+| [`IfExists`](https://pkg.go.dev/github.com/bitfield/script#IfExists) | do something only if some file exists |
+| [`ListFiles`](https://pkg.go.dev/github.com/bitfield/script#ListFiles) | file listing (including wildcards) |
+| [`Post`](https://pkg.go.dev/github.com/bitfield/script#Post) | HTTP response |
+| [`Slice`](https://pkg.go.dev/github.com/bitfield/script#Slice) | slice elements, one per line |
+| [`Stdin`](https://pkg.go.dev/github.com/bitfield/script#Stdin) | standard input |
+
+## Modifiers
+
+These are methods on a pipe that change its configuration:
+
+| Source | Modifies |
+| -------- | ------------- |
+| [`WithEnv`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithEnv) | environment for commands |
+| [`WithError`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithError) | pipe error status |
+| [`WithHTTPClient`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithHTTPClient) | client for HTTP requests |
+| [`WithReader`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithReader) | pipe source |
+| [`WithStderr`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithStderr) | standard error output stream for command |
+| [`WithStdout`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithStdout) | standard output stream for pipe |
 
 ## Filters
 
@@ -340,7 +353,8 @@ Sinks are methods that return some data from a pipe, ending the pipeline and ext
 
 | Version | New |
 | ----------- | ------- |
-| _next_  | [`DecodeBase64`](https://pkg.go.dev/github.com/bitfield/script#Pipe.DecodeBase64) / [`EncodeBase64`](https://pkg.go.dev/github.com/bitfield/script#Pipe.EncodeBase64) |
+| _next_  | [`WithEnv`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithEnv) |
+|         | [`DecodeBase64`](https://pkg.go.dev/github.com/bitfield/script#Pipe.DecodeBase64) / [`EncodeBase64`](https://pkg.go.dev/github.com/bitfield/script#Pipe.EncodeBase64) |
 | v0.22.0 | [`Tee`](https://pkg.go.dev/github.com/bitfield/script#Pipe.Tee), [`WithStderr`](https://pkg.go.dev/github.com/bitfield/script#Pipe.WithStderr) |
 | v0.21.0 | HTTP support: [`Do`](https://pkg.go.dev/github.com/bitfield/script#Pipe.Do), [`Get`](https://pkg.go.dev/github.com/bitfield/script#Pipe.Get), [`Post`](https://pkg.go.dev/github.com/bitfield/script#Pipe.Post) |
 | v0.20.0 | [`JQ`](https://pkg.go.dev/github.com/bitfield/script#Pipe.JQ) |

script.go

@@ -32,10 +32,10 @@ type Pipe struct {
 	stdout     io.Writer
 	httpClient *http.Client
 
-	// because pipe stages are concurrent, protect 'err' and 'stderr'
 	mu     *sync.Mutex
 	err    error
 	stderr io.Writer
+	env    []string
 }
 
 // Args creates a pipe containing the program's command-line arguments from
@@ -168,6 +168,7 @@ func NewPipe() *Pipe {
 		mu:         new(sync.Mutex),
 		stdout:     os.Stdout,
 		httpClient: http.DefaultClient,
+		env:        nil,
 	}
 }
 
@@ -374,6 +375,12 @@ func (p *Pipe) EncodeBase64() *Pipe {
 	})
 }
 
+func (p *Pipe) environment() []string {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	return p.env
+}
+
 // Error returns any error present on the pipe, or nil otherwise.
 // Error is not a sink and does not wait until the pipe reaches
 // completion. To wait for completion before returning the error,
@@ -392,6 +399,11 @@ func (p *Pipe) Error() error {
 // error output). The effect of this is to filter the contents of the pipe
 // through the external command.
 //
+// # Environment
+//
+// The command inherits the current process's environment, optionally modified
+// by [Pipe.WithEnv].
+//
 // # Error handling
 //
 // If the command had a non-zero exit status, the pipe's error status will also
@@ -419,6 +431,10 @@ func (p *Pipe) Exec(cmdLine string) *Pipe {
 		if pipeStderr != nil {
 			cmd.Stderr = pipeStderr
 		}
+		pipeEnv := p.environment()
+		if pipeEnv != nil {
+			cmd.Env = pipeEnv
+		}
 		err = cmd.Start()
 		if err != nil {
 			fmt.Fprintln(cmd.Stderr, err)
@@ -430,7 +446,8 @@ func (p *Pipe) Exec(cmdLine string) *Pipe {
 
 // ExecForEach renders cmdLine as a Go template for each line of input, running
 // the resulting command, and produces the combined output of all these
-// commands in sequence. See [Pipe.Exec] for error handling details.
+// commands in sequence. See [Pipe.Exec] for details on error handling and
+// environment variables.
 //
 // This is mostly useful for substituting data into commands using Go template
 // syntax. For example:
@@ -460,6 +477,9 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe {
 			if pipeStderr != nil {
 				cmd.Stderr = pipeStderr
 			}
+			if p.env != nil {
+				cmd.Env = p.env
+			}
 			err = cmd.Start()
 			if err != nil {
 				fmt.Fprintln(cmd.Stderr, err)
@@ -903,6 +923,16 @@ func (p *Pipe) Wait() error {
 	return p.Error()
 }
 
+// WithEnv sets the environment for subsequent [Pipe.Exec] and [Pipe.ExecForEach]
+// commands to the string slice env, using the same format as [os/exec.Cmd.Env].
+// An empty slice unsets all existing environment variables.
+func (p *Pipe) WithEnv(env []string) *Pipe {
+	p.mu.Lock()
+	defer p.mu.Unlock()
+	p.env = env
+	return p
+}
+
 // WithError sets the error err on the pipe.
 func (p *Pipe) WithError(err error) *Pipe {
 	p.SetError(err)

script_test.go

@@ -1768,6 +1768,40 @@ func TestWithStdout_SetsSpecifiedWriterAsStdout(t *testing.T) {
 	}
 }
 
+func TestWithEnv_UnsetsAllEnvVarsGivenEmptySlice(t *testing.T) {
+	t.Parallel()
+	p := script.NewPipe().WithEnv([]string{"ENV1=test1"}).Exec("sh -c 'echo ENV1=$ENV1'")
+	want := "ENV1=test1\n"
+	got, err := p.String()
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got != want {
+		t.Fatalf("want %q, got %q", want, got)
+	}
+	got, err = p.Echo("").WithEnv([]string{}).Exec("sh -c 'echo ENV1=$ENV1'").String()
+	if err != nil {
+		t.Fatal(err)
+	}
+	want = "ENV1=\n"
+	if got != want {
+		t.Errorf("want %q, got %q", want, got)
+	}
+}
+
+func TestWithEnv_SetsGivenVariablesForSubsequentExec(t *testing.T) {
+	t.Parallel()
+	env := []string{"ENV1=test1", "ENV2=test2"}
+	got, err := script.NewPipe().WithEnv(env).Exec("sh -c 'echo ENV1=$ENV1 ENV2=$ENV2'").String()
+	if err != nil {
+		t.Fatal(err)
+	}
+	want := "ENV1=test1 ENV2=test2\n"
+	if got != want {
+		t.Errorf("want %q, got %q", want, got)
+	}
+}
+
 func TestErrorReturnsErrorSetByPreviousPipeStage(t *testing.T) {
 	t.Parallel()
 	p := script.File("testdata/nonexistent.txt")