sh.go

Documentation: github.com/codeskyblue/go-sh

     1  /*
     2  Package go-sh is intended to make shell call with golang more easily.
     3  Some usage is more similar to os/exec, eg: Run(), Output(), Command(name, args...)
     4  
     5  But with these similar function, pipe is added in and this package also got shell-session support.
     6  
     7  Why I love golang so much, because the usage of golang is simple, but the power is unlimited. I want to make this pakcage got the sample style like golang.
     8  
     9  	// just like os/exec
    10  	sh.Command("echo", "hello").Run()
    11  
    12  	// support pipe
    13  	sh.Command("echo", "hello").Command("wc", "-c").Run()
    14  
    15  	// create a session to store dir and env
    16  	sh.NewSession().SetDir("/").Command("pwd")
    17  
    18  	// shell buildin command - "test"
    19  	sh.Test("dir", "mydir")
    20  
    21  	// like shell call: (cd /; pwd)
    22  	sh.Command("pwd", sh.Dir("/")) same with sh.Command(sh.Dir("/"), "pwd")
    23  
    24  	// output to json and xml easily
    25  	v := map[string] int {}
    26  	err = sh.Command("echo", `{"number": 1}`).UnmarshalJSON(&v)
    27  */
    28  package sh
    29  
    30  import (
    31  	"fmt"
    32  	"io"
    33  	"os"
    34  	"os/exec"
    35  	"reflect"
    36  	"strings"
    37  	"time"
    38  
    39  	"github.com/codegangsta/inject"
    40  )
    41  
    42  type Dir string
    43  
    44  type Session struct {
    45  	inj     inject.Injector
    46  	alias   map[string][]string
    47  	cmds    []*exec.Cmd
    48  	dir     Dir
    49  	started bool
    50  	Env     map[string]string
    51  	Stdin   io.Reader
    52  	Stdout  io.Writer
    53  	Stderr  io.Writer
    54  	ShowCMD bool // enable for debug
    55  	timeout time.Duration
    56  
    57  	// additional pipe options
    58  	PipeFail      bool // returns error of rightmost no-zero command
    59  	PipeStdErrors bool // combine std errors of all pipe commands
    60  }
    61  
    62  func (s *Session) writePrompt(args ...interface{}) {
    63  	var ps1 = fmt.Sprintf("[golang-sh]$")
    64  	args = append([]interface{}{ps1}, args...)
    65  	fmt.Fprintln(s.Stderr, args...)
    66  }
    67  
    68  func NewSession() *Session {
    69  	env := make(map[string]string)
    70  	for _, key := range []string{"PATH"} {
    71  		env[key] = os.Getenv(key)
    72  	}
    73  	s := &Session{
    74  		inj:    inject.New(),
    75  		alias:  make(map[string][]string),
    76  		dir:    Dir(""),
    77  		Stdin:  strings.NewReader(""),
    78  		Stdout: os.Stdout,
    79  		Stderr: os.Stderr,
    80  		Env:    env,
    81  	}
    82  	return s
    83  }
    84  
    85  func InteractiveSession() *Session {
    86  	s := NewSession()
    87  	s.SetStdin(os.Stdin)
    88  	return s
    89  }
    90  
    91  func Command(name string, a ...interface{}) *Session {
    92  	s := NewSession()
    93  	return s.Command(name, a...)
    94  }
    95  
    96  func Echo(in string) *Session {
    97  	s := NewSession()
    98  	return s.SetInput(in)
    99  }
   100  
   101  func (s *Session) Alias(alias, cmd string, args ...string) {
   102  	v := []string{cmd}
   103  	v = append(v, args...)
   104  	s.alias[alias] = v
   105  }
   106  
   107  func (s *Session) Command(name string, a ...interface{}) *Session {
   108  	var args = make([]string, 0)
   109  	var sType = reflect.TypeOf("")
   110  
   111  	// init cmd, args, dir, envs
   112  	// if not init, program may panic
   113  	s.inj.Map(name).Map(args).Map(s.dir).Map(map[string]string{})
   114  	for _, v := range a {
   115  		switch reflect.TypeOf(v) {
   116  		case sType:
   117  			args = append(args, v.(string))
   118  		default:
   119  			s.inj.Map(v)
   120  		}
   121  	}
   122  	if len(args) != 0 {
   123  		s.inj.Map(args)
   124  	}
   125  	s.inj.Invoke(s.appendCmd)
   126  	return s
   127  }
   128  
   129  // combine Command and Run
   130  func (s *Session) Call(name string, a ...interface{}) error {
   131  	return s.Command(name, a...).Run()
   132  }
   133  
   134  /*
   135  func (s *Session) Exec(cmd string, args ...string) error {
   136  	return s.Call(cmd, args)
   137  }
   138  */
   139  
   140  func (s *Session) SetEnv(key, value string) *Session {
   141  	s.Env[key] = value
   142  	return s
   143  }
   144  
   145  func (s *Session) SetDir(dir string) *Session {
   146  	s.dir = Dir(dir)
   147  	return s
   148  }
   149  
   150  func (s *Session) SetInput(in string) *Session {
   151  	s.Stdin = strings.NewReader(in)
   152  	return s
   153  }
   154  
   155  func (s *Session) SetStdin(r io.Reader) *Session {
   156  	s.Stdin = r
   157  	return s
   158  }
   159  
   160  func (s *Session) SetTimeout(d time.Duration) *Session {
   161  	s.timeout = d
   162  	return s
   163  }
   164  
   165  func newEnviron(env map[string]string, inherit bool) []string { //map[string]string {
   166  	environ := make([]string, 0, len(env))
   167  	if inherit {
   168  		for _, line := range os.Environ() {
   169  			for k, _ := range env {
   170  				if strings.HasPrefix(line, k+"=") {
   171  					goto CONTINUE
   172  				}
   173  			}
   174  			environ = append(environ, line)
   175  		CONTINUE:
   176  		}
   177  	}
   178  	for k, v := range env {
   179  		environ = append(environ, k+"="+v)
   180  	}
   181  	return environ
   182  }
   183  
   184  func (s *Session) appendCmd(cmd string, args []string, cwd Dir, env map[string]string) {
   185  	if s.started {
   186  		s.started = false
   187  		s.cmds = make([]*exec.Cmd, 0)
   188  	}
   189  	for k, v := range s.Env {
   190  		if _, ok := env[k]; !ok {
   191  			env[k] = v
   192  		}
   193  	}
   194  	environ := newEnviron(s.Env, true) // true: inherit sys-env
   195  	v, ok := s.alias[cmd]
   196  	if ok {
   197  		cmd = v[0]
   198  		args = append(v[1:], args...)
   199  	}
   200  	c := exec.Command(cmd, args...)
   201  	c.Env = environ
   202  	c.Dir = string(cwd)
   203  	s.cmds = append(s.cmds, c)
   204  }
   205
View as plain text