feat: exec with context

script.go

@@ -3,6 +3,7 @@ package script
 import (
 	"bufio"
 	"container/ring"
+	"context"
 	"crypto/sha256"
 	"encoding/base64"
 	"encoding/hex"
@@ -32,6 +33,7 @@ type Pipe struct {
 	stdout     io.Writer
 	httpClient *http.Client
 
+	ctx    context.Context
 	mu     *sync.Mutex
 	err    error
 	stderr io.Writer
@@ -166,6 +168,7 @@ func NewPipe() *Pipe {
 	return &Pipe{
 		Reader:     ReadAutoCloser{},
 		mu:         new(sync.Mutex),
+		ctx:        context.Background(),
 		stdout:     os.Stdout,
 		httpClient: http.DefaultClient,
 		env:        nil,
@@ -423,7 +426,7 @@ func (p *Pipe) Exec(cmdLine string) *Pipe {
 		if err != nil {
 			return err
 		}
-		cmd := exec.Command(args[0], args[1:]...)
+		cmd := exec.CommandContext(p.ctx, args[0], args[1:]...)
 		cmd.Stdin = r
 		cmd.Stdout = w
 		cmd.Stderr = w
@@ -470,7 +473,7 @@ func (p *Pipe) ExecForEach(cmdLine string) *Pipe {
 			if err != nil {
 				return err
 			}
-			cmd := exec.Command(args[0], args[1:]...)
+			cmd := exec.CommandContext(p.ctx, args[0], args[1:]...)
 			cmd.Stdout = w
 			cmd.Stderr = w
 			pipeStderr := p.stdErr()
@@ -974,6 +977,13 @@ func (p *Pipe) WithStdout(w io.Writer) *Pipe {
 	return p
 }
 
+// WithContext sets context.Context for the pipe. Adds support for graceful pipe
+// shutdown. Currently works with [Pipe.Exec] and [Pipe.ExecForEach]
+func (p *Pipe) WithContext(ctx context.Context) *Pipe {
+	p.ctx = ctx
+	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

@@ -3,6 +3,7 @@ package script_test
 import (
 	"bufio"
 	"bytes"
+	"context"
 	"errors"
 	"fmt"
 	"io"
@@ -15,6 +16,7 @@ import (
 	"strings"
 	"testing"
 	"testing/iotest"
+	"time"
 
 	"github.com/bitfield/script"
 	"github.com/google/go-cmp/cmp"
@@ -1219,6 +1221,48 @@ func TestExecRunsGoHelpAndGetsUsageMessage(t *testing.T) {
 	}
 }
 
+func TestWithContextTimeout(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithTimeout(context.Background(), 1*time.Second)
+	defer cancel()
+	p := script.NewPipe().WithContext(ctx).Exec("go run ./testdata/test_cli.go sleep 10")
+	p.Wait()
+	err := p.Error()
+	if err != nil && err.Error() != "signal: killed" {
+		t.Fatalf("context should timeout, %v", err)
+	}
+	t.Log(p.ExitStatus())
+}
+
+func TestWithContextTimeoutBeforeRun(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithTimeout(context.Background(), 0*time.Second)
+	defer cancel()
+	p := script.NewPipe().WithContext(ctx).Exec("go run ./testdata/test_cli.go sleep 10")
+	p.Wait()
+	err := p.Error()
+	if !errors.Is(err, context.DeadlineExceeded) {
+		t.Fatal("context should timeout")
+	}
+	t.Log(p.ExitStatus())
+}
+
+func TestWithContextCancel(t *testing.T) {
+	t.Parallel()
+	ctx, cancel := context.WithCancel(context.Background())
+	p := script.NewPipe().WithContext(ctx).Exec("go run ./testdata/test_cli.go sleep 2")
+	go func() {
+		<-time.After(1 * time.Second)
+		cancel()
+	}()
+	p.Wait()
+	err := p.Error()
+	if errors.Is(err, context.Canceled) {
+		t.Fatalf("context should cancel")
+	}
+	t.Log(p.ExitStatus())
+}
+
 func TestFileOutputsContentsOfSpecifiedFile(t *testing.T) {
 	t.Parallel()
 	want := "This is the first line in the file.\nHello, world.\nThis is another line in the file.\n"

testdata/test_cli.go

@@ -0,0 +1,67 @@
+package main
+
+import (
+	"fmt"
+	"os"
+	"os/signal"
+	"strconv"
+	"strings"
+	"syscall"
+	"time"
+)
+
+const commandSleep = "sleep"
+
+func main() {
+	if len(os.Args) < 2 {
+		fmt.Println("No command provided")
+		usage()
+		os.Exit(1)
+	}
+
+	handleSignals()
+
+	switch strings.ToLower(os.Args[1]) {
+	case commandSleep:
+		if len(os.Args) != 3 {
+			fmt.Printf("Usage: %s %s <seconds>\n", os.Args[0], commandSleep)
+			usage()
+			os.Exit(1)
+		}
+		err := sleep(os.Args[2])
+		if err != nil {
+			fmt.Printf("Error sleeping: %s\n", err)
+			os.Exit(1)
+		}
+	default:
+		fmt.Printf("Unknown command: %s\n", os.Args[1])
+		usage()
+		os.Exit(1)
+	}
+}
+
+func usage() {
+	fmt.Printf("Usage of %s:\n", os.Args[0])
+	fmt.Printf("     sleep <seconds>\n")
+}
+
+func sleep(seconds string) error {
+	s, err := strconv.Atoi(seconds)
+	if err != nil {
+		return fmt.Errorf("sleep expects an integer, got %s", seconds)
+	}
+	<-time.After(time.Duration(s) * time.Second)
+	return nil
+}
+
+func handleSignals() {
+	sigs := make(chan os.Signal, 1)
+
+	signal.Notify(sigs, syscall.SIGINT)
+
+	go func() {
+		<-sigs
+		fmt.Println("\nReceived Ctrl+C, exiting...")
+		os.Exit(0)
+	}()
+}