sh.go

package main

import (
	"bytes"
	"fmt"
	"github.com./hashicorp/go-hclog"
	"io"
	"log"
	"os"
	"os/exec"
	"strings"
)

type Module struct {
	logger hclog.Logger
}

func New(logger hclog.Logger) *Module {
	return &Module{logger: logger}
}

type ExecutionError struct {
	Command  string
	Args     []string
	ExitCode int
}

func (t ExecutionError) Error() string {
	return fmt.Sprintf(`running "%s %s" failed with exit code %d`, t.Command, strings.Join(t.Args, " "), t.ExitCode)
}

// RunCmd returns a function that will call Run with the given command. This is
// 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 :(
//
//  // 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")
//
// RunCmd uses Exec underneath, so see those docs for more details.
func (sh *Module) RunCmd(cmd string, args ...string) func(args ...string) error {
	return func(args2 ...string) error {
		return sh.Run(cmd, append(args, args2...)...)
	}
}

// OutCmd is like RunCmd except the command returns the output of the
// command.
func (sh *Module) OutCmd(cmd string, args ...string) func(args ...string) (string, error) {
	return func(args2 ...string) (string, error) {
		return sh.Output(cmd, append(args, args2...)...)
	}
}

// Run is like RunWith, but doesn't specify any environment variables.
func (sh *Module) Run(cmd string, args ...string) error {
	return sh.RunWith(nil, cmd, args...)
}

// RunV is like Run, but always sends the command's stdout to os.Stdout.
func (sh *Module) RunV(cmd string, args ...string) error {
	loggerOpts := &hclog.StandardLoggerOptions{}
	stdLogger := sh.logger.StandardWriter(loggerOpts)
	_, err := sh.Exec(nil, stdLogger, stdLogger, cmd, args...)
	return err
}

// RunWith runs the given command, directing stderr to this program's stderr and
// printing stdout to stdout if mage was run with -v.  It adds adds env to the
// environment variables for the command being run. Environment variables should
// be in the format name=value.
func (sh *Module) RunWith(env map[string]string, cmd string, args ...string) error {
	var output io.Writer
	loggerOpts := &hclog.StandardLoggerOptions{}
	stdLogger := sh.logger.StandardWriter(loggerOpts)
	if sh.logger.IsDebug() {
		output = stdLogger
	}
	_, err := sh.Exec(env, output, stdLogger, cmd, args...)
	return err
}

// RunWithV is like RunWith, but always sends the command's stdout to os.Stdout.
func (sh *Module) RunWithV(env map[string]string, cmd string, args ...string) error {
	loggerOpts := &hclog.StandardLoggerOptions{}
	stdLogger := sh.logger.StandardWriter(loggerOpts)
	_, err := sh.Exec(env, stdLogger, stdLogger, cmd, args...)
	return err
}

// Output runs the command and returns the text from stdout.
func (sh *Module) Output(cmd string, args ...string) (string, error) {
	loggerOpts := &hclog.StandardLoggerOptions{}
	stdLogger := sh.logger.StandardWriter(loggerOpts)
	buf := &bytes.Buffer{}
	_, err := sh.Exec(nil, buf, stdLogger, cmd, args...)
	return strings.TrimSuffix(buf.String(), "\n"), err
}

// OutputWith is like RunWith, but returns what is written to stdout.
func (sh *Module) OutputWith(env map[string]string, cmd string, args ...string) (string, error) {
	loggerOpts := &hclog.StandardLoggerOptions{}
	stdlogger := sh.logger.StandardWriter(loggerOpts)
	buf := &bytes.Buffer{}
	_, err := sh.Exec(env, buf, stdlogger, cmd, args...)
	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.
//
// 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 (sh *Module) Exec(env map[string]string, stdout, stderr io.Writer, cmd string, args ...string) (ran bool, err error) {
	expand := func(s string) string {
		s2, ok := 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)
	}
	ran, code, err := sh.run(env, stdout, stderr, cmd, args...)
	if err == nil {
		return true, nil
	}
	if ran {
		return ran, ExecutionError{cmd, args, code}
	}
	return ran, fmt.Errorf(`failed to run "%s %s: %v"`, cmd, strings.Join(args, " "), err)
}

func (sh *Module) 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)
	}
	c.Stderr = stderr
	c.Stdout = stdout
	c.Stdin = os.Stdin
	log.Println("exec:", cmd, strings.Join(args, " "))
	err = c.Run()
	return sh.CmdRan(err), sh.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
// error is an unrecognized type, or it is an error from exec.Command that says
// the command failed to run (usually due to the command not existing or not
// being executable), it reports false.
func (sh *Module) CmdRan(err error) bool {
	if err == nil {
		return true
	}
	ee, ok := err.(*exec.ExitError)
	if ok {
		return ee.Exited()
	}
	_, ok = err.(ExecutionError)
	if ok {
		return true
	}
	return false
}

type exitStatus interface {
	ExitStatus() int
}

// ExitStatus returns the exit status of the error if it is an exec.ExitError
// or if it implements ExitStatus() int.
// 0 if it is nil or 1 if it is a different error.
func (sh *Module) ExitStatus(err error) int {
	if err == nil {
		return 0
	}
	if e, ok := err.(ExecutionError); ok {
		return e.ExitCode
	}
	if e, ok := err.(exitStatus); ok {
		return e.ExitStatus()
	}
	if e, ok := err.(*exec.ExitError); ok {
		if ex, ok := e.Sys().(exitStatus); ok {
			return ex.ExitStatus()
		}
	}
	return 1
}