github.com/NeowayLabs/nash@v0.2.2-0.20200127205349-a227041ffd50/internal/sh/builtin/split.go (about)

     1  package builtin
     2  
     3  import (
     4  	"io"
     5  	"strings"
     6  
     7  	"github.com/madlambda/nash/errors"
     8  	"github.com/madlambda/nash/sh"
     9  )
    10  
    11  type (
    12  	splitFn struct {
    13  		content string
    14  		sep     sh.Obj
    15  	}
    16  )
    17  
    18  func newSplit() *splitFn {
    19  	return &splitFn{}
    20  }
    21  
    22  func (s *splitFn) ArgNames() []sh.FnArg {
    23  	return []sh.FnArg{
    24  		sh.NewFnArg("sep", false),
    25  		sh.NewFnArg("content", false),
    26  	}
    27  }
    28  
    29  func (s *splitFn) Run(in io.Reader, out io.Writer, err io.Writer) ([]sh.Obj, error) {
    30  	var output []string
    31  
    32  	content := s.content
    33  
    34  	switch s.sep.Type() {
    35  	case sh.StringType:
    36  		sep := s.sep.(*sh.StrObj).Str()
    37  		output = strings.Split(content, sep)
    38  	case sh.ListType:
    39  		sepList := s.sep.(*sh.ListObj).List()
    40  		output = splitByList(content, sepList)
    41  	case sh.FnType:
    42  		sepFn := s.sep.(*sh.FnObj).Fn()
    43  		output = splitByFn(content, sepFn)
    44  	default:
    45  		return nil, errors.NewError("Invalid separator value: %v", s.sep)
    46  	}
    47  
    48  	listobjs := make([]sh.Obj, len(output))
    49  	for i := 0; i < len(output); i++ {
    50  		listobjs[i] = sh.NewStrObj(output[i])
    51  	}
    52  
    53  	return []sh.Obj{sh.NewListObj(listobjs)}, nil
    54  }
    55  
    56  func (s *splitFn) SetArgs(args []sh.Obj) error {
    57  	if len(args) != 2 {
    58  		return errors.NewError("split: expects 2 parameters")
    59  	}
    60  
    61  	if args[0].Type() != sh.StringType {
    62  		return errors.NewError("split: first parameter must be a string")
    63  	}
    64  
    65  	content := args[0].(*sh.StrObj)
    66  	s.content = content.Str()
    67  	s.sep = args[1]
    68  	return nil
    69  }
    70  
    71  func splitByList(content string, delims []sh.Obj) []string {
    72  	return strings.FieldsFunc(content, func(r rune) bool {
    73  		for _, delim := range delims {
    74  			if delim.Type() != sh.StringType {
    75  				continue
    76  			}
    77  
    78  			objstr := delim.(*sh.StrObj)
    79  
    80  			if len(objstr.Str()) > 0 && rune(objstr.Str()[0]) == r {
    81  				return true
    82  			}
    83  		}
    84  
    85  		return false
    86  	})
    87  }
    88  
    89  func splitByFn(content string, splitFunc sh.FnDef) []string {
    90  	return strings.FieldsFunc(content, func(r rune) bool {
    91  		fn := splitFunc.Build()
    92  		arg := sh.NewStrObj(string(r))
    93  		fn.SetArgs([]sh.Obj{arg})
    94  		err := fn.Start()
    95  		if err != nil {
    96  			return false
    97  		}
    98  
    99  		err = fn.Wait()
   100  		if err != nil {
   101  			return false
   102  		}
   103  
   104  		results := fn.Results()
   105  		if len(results) != 1 {
   106  			// expects a single return fn
   107  			return false
   108  		}
   109  
   110  		result := results[0]
   111  
   112  		//FIXME: It would be cool to only accept booleans
   113  		// since the splitter is a predicate
   114  		if result.Type() != sh.StringType {
   115  			return false
   116  		}
   117  
   118  		status := result.(*sh.StrObj)
   119  		if status.Str() == "0" {
   120  			return true
   121  		}
   122  
   123  		return false
   124  	})
   125  }