Added working directory and context support.

sh/cmd.go

@@ -2,6 +2,7 @@ package sh
 
 import (
 	"bytes"
+	"context"
 	"fmt"
 	"io"
 	"log"
@@ -16,19 +17,19 @@ import (
 // useful for creating command aliases to make your scripts easier to read, like
 // this:
 //
-//  // in a helper file somewhere
-//  var g0 = sh.RunCmd("go")  // go is a keyword :(
+//	 // in a helper file somewhere
+//	 var g0 = sh.RunCmd("go")  // go is a keyword :(
 //
-//  // somewhere in your main code
-//	if err := g0("install", "github.com/gohugo/hugo"); err != nil {
-//		return err
-//  }
+//	 // somewhere in your main code
+//		if err := g0("install", "github.com/gohugo/hugo"); err != nil {
+//			return err
+//	 }
 //
 // Args passed to command get baked in as args to the command when you run it.
 // Any args passed in when you run the returned function will be appended to the
 // original args.  For example, this is equivalent to the above:
 //
-//  var goInstall = sh.RunCmd("go", "install") goInstall("github.com/gohugo/hugo")
+//	var goInstall = sh.RunCmd("go", "install") goInstall("github.com/gohugo/hugo")
 //
 // RunCmd uses Exec underneath, so see those docs for more details.
 func RunCmd(cmd string, args ...string) func(args ...string) error {
@@ -89,53 +90,123 @@ func OutputWith(env map[string]string, cmd string, args ...string) (string, erro
 	return strings.TrimSuffix(buf.String(), "\n"), err
 }
 
-// Exec executes the command, piping its stderr to mage's stderr and
-// piping its stdout to the given writer. If the command fails, it will return
-// an error that, if returned from a target or mg.Deps call, will cause mage to
-// exit with the same code as the command failed with.  Env is a list of
-// environment variables to set when running the command, these override the
-// current environment variables set (which are also passed to the command). cmd
-// and args may include references to environment variables in $FOO format, in
-// which case these will be expanded before the command is run.
+// Exec executes the command, piping its stdout and stderr to the given writers.
+// If the command fails, it will return an error that, if returned from a target
+// or mg.Deps call, will cause mage to exit with the same code as the command
+// failed with.
+// Env is a list of environment variables to set when running the command,
+// these override the current environment variables set (which are also passed
+// to the command).
+// cmd and args may include references to environment variables in $FOO format,
+// in which case these will be expanded before the command is run.
 //
 // Ran reports if the command ran (rather than was not found or not executable).
 // Code reports the exit code the command returned if it ran. If err == nil, ran
 // is always true and code is always 0.
 func Exec(env map[string]string, stdout, stderr io.Writer, cmd string, args ...string) (ran bool, err error) {
+	return Command{
+		Cmd:    cmd,
+		Args:   args,
+		Stdout: stdout,
+		Stderr: stderr,
+		Env:    env,
+	}.Exec(context.Background())
+}
+
+// Command is a command to be executed.
+//
+// Both Path and Args may include references to environment variables in $FOO
+// format, // in which case these will be expanded before the command is run.
+type Command struct {
+	// Cmd is the path of the command to execute.
+	// Relative paths are evaluated with respect to WorkingDir.
+	//
+	// Environment variable references of the form $FOO will be expanded before
+	// the command is run.
+	Cmd string
+	// Args are the command line arguments to pass to the command.
+	//
+	// Environment variable references of the form $FOO will be expanded before
+	// the command is run.
+	Args []string
+
+	// Env is a list of environment variables to set when running the command.
+	// These override the current environment variables set (which are also
+	// passed to the command).
+	Env map[string]string
+
+	// Stdout is the command's stdout stream.
+	Stdout io.Writer
+	// Stderr is the command's stderr stream.
+	Stderr io.Writer
+
+	// WorkingDir specifies the working directory this will command execute in.
+	// An empty string indicates the command should run in the current working
+	// directory.
+	WorkingDir string
+}
+
+// Output and Exec use value receivers to avoid race conditions when modifying
+// the Command's internal state during shell expansion.
+
+// Output runs the [Command] and returns the text from stdout.
+//
+// See [Command.Exec] for more detailts.
+func (cmd Command) Output(ctx context.Context) (string, error) {
+	buf := &bytes.Buffer{}
+	cmd.Stdout = buf
+	_, err := cmd.Exec(ctx)
+	return strings.TrimSuffix(buf.String(), "\n"), err
+}
+
+// Exec executes the [Command] using the provided [context.Context] for
+// cancellation, and piping the stdout and stderr to the given writers.
+// If the command fails, it will return an error that, if returned from a target
+// or [mg.Deps] call, will cause mage to exit with the same code as the command
+// failed with.
+//
+// Ran reports if the Command ran (rather than was not found or not executable).
+// Code reports the exit code the command returned if it ran, and can be
+// retrieved from err with [mg.ExitStatus].
+// If err == nil, ran is always true and code is always 0.
+func (cmd Command) Exec(ctx context.Context) (ran bool, err error) {
 	expand := func(s string) string {
-		s2, ok := env[s]
+		s2, ok := cmd.Env[s]
 		if ok {
 			return s2
 		}
 		return os.Getenv(s)
 	}
-	cmd = os.Expand(cmd, expand)
-	for i := range args {
-		args[i] = os.Expand(args[i], expand)
+	cmd.Cmd = os.Expand(cmd.Cmd, expand)
+	for i := range cmd.Args {
+		cmd.Args[i] = os.Expand(cmd.Args[i], expand)
 	}
-	ran, code, err := run(env, stdout, stderr, cmd, args...)
+	ran, code, err := cmd.run(ctx)
 	if err == nil {
 		return true, nil
 	}
 	if ran {
-		return ran, mg.Fatalf(code, `running "%s %s" failed with exit code %d`, cmd, strings.Join(args, " "), code)
+		return ran, mg.Fatalf(code, `running "%s %s" failed with exit code %d`, cmd.Cmd, strings.Join(cmd.Args, " "), code)
 	}
-	return ran, fmt.Errorf(`failed to run "%s %s: %v"`, cmd, strings.Join(args, " "), err)
+	return ran, fmt.Errorf(`failed to run "%s %s: %v"`, cmd.Cmd, strings.Join(cmd.Args, " "), err)
+
 }
 
-func run(env map[string]string, stdout, stderr io.Writer, cmd string, args ...string) (ran bool, code int, err error) {
-	c := exec.Command(cmd, args...)
-	c.Env = os.Environ()
-	for k, v := range env {
-		c.Env = append(c.Env, k+"="+v)
+func (cmd *Command) run(ctx context.Context) (ran bool, code int, err error) {
+	c := exec.CommandContext(ctx, cmd.Cmd, cmd.Args...)
+	env := os.Environ()
+	for k, v := range cmd.Env {
+		env = append(c.Env, k+"="+v)
 	}
-	c.Stderr = stderr
-	c.Stdout = stdout
+	c.Env = env
+	c.Stderr = cmd.Stderr
+	c.Stdout = cmd.Stdout
 	c.Stdin = os.Stdin
+	c.Dir = cmd.WorkingDir
 
-	var quoted []string 
-	for i := range args {
-		quoted = append(quoted, fmt.Sprintf("%q", args[i]));
+	quoted := make([]string, 0, len(cmd.Args))
+	for _, c := range cmd.Args {
+		quoted = append(quoted, fmt.Sprintf("%q", c))
 	}
 	// To protect against logging from doing exec in global variables
 	if mg.Verbose() {
@@ -144,6 +215,7 @@ func run(env map[string]string, stdout, stderr io.Writer, cmd string, args ...st
 	err = c.Run()
 	return CmdRan(err), ExitStatus(err), err
 }
+
 // CmdRan examines the error to determine if it was generated as a result of a
 // command running via os/exec.Command.  If the error is nil, or the command ran
 // (even if it exited with a non-zero exit code), CmdRan reports true.  If the

sh/cmd_test.go

@@ -2,8 +2,10 @@ package sh
 
 import (
 	"bytes"
+	"context"
 	"os"
 	"testing"
+	"time"
 )
 
 func TestOutCmd(t *testing.T) {
@@ -68,5 +70,40 @@ func TestAutoExpand(t *testing.T) {
 	if s != "baz" {
 		t.Fatalf(`Expected "baz" but got %q`, s)
 	}
+}
 
+func TestContextTimeout(t *testing.T) {
+	d := 1 * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), d)
+	defer cancel()
+	start := time.Now()
+	_, err := Command{
+		Cmd:  os.Args[0],
+		Args: []string{"-sleep", (2 * d).String()},
+	}.Exec(ctx)
+	dd := time.Since(start)
+	if err == nil {
+		t.Fatalf("Command should have errored")
+	}
+	if dd < d {
+		t.Fatalf("Duration too short: expected %v, got %v", d, dd)
+	}
+	// allow some wiggle room, too account for Exec overheard
+	if dd-d > 50*time.Millisecond {
+		t.Fatalf("Expected duration %v, got %v", d, dd)
+	}
+}
+
+func TestWorkingDirectory(t *testing.T) {
+	tmp := t.TempDir()
+	s, err := Command{
+		Cmd:        "pwd",
+		WorkingDir: tmp,
+	}.Output(context.Background())
+	if err != nil {
+		t.Fatal(err)
+	}
+	if s != tmp {
+		t.Fatalf(`Expected %q but got %q`, tmp, s)
+	}
 }

sh/testmain_test.go

@@ -5,6 +5,7 @@ import (
 	"fmt"
 	"os"
 	"testing"
+	"time"
 )
 
 var (
@@ -14,6 +15,7 @@ var (
 	stdout    string
 	exitCode  int
 	printVar  string
+	sleep     time.Duration
 )
 
 func init() {
@@ -23,6 +25,7 @@ func init() {
 	flag.StringVar(&stdout, "stdout", "", "")
 	flag.IntVar(&exitCode, "exit", 0, "")
 	flag.StringVar(&printVar, "printVar", "", "")
+	flag.DurationVar(&sleep, "sleep", 0, "")
 }
 
 func TestMain(m *testing.M) {
@@ -37,6 +40,11 @@ func TestMain(m *testing.M) {
 		return
 	}
 
+	if sleep != 0 {
+		time.Sleep(sleep)
+		return
+	}
+
 	if helperCmd {
 		fmt.Fprintln(os.Stderr, stderr)
 		fmt.Fprintln(os.Stdout, stdout)