github.com/kolbycrouch/elvish@v0.14.1-0.20210614162631-215b9ac1c423/pkg/eval/mods/re/re.go (about)

     1  // Package re implements a regular expression module.
     2  package re
     3  
     4  import (
     5  	"fmt"
     6  	"regexp"
     7  
     8  	"src.elv.sh/pkg/eval"
     9  	"src.elv.sh/pkg/eval/vals"
    10  	"src.elv.sh/pkg/persistent/vector"
    11  )
    12  
    13  // Ns is the namespace for the re: module.
    14  var Ns = eval.NsBuilder{}.AddGoFns("re:", fns).Ns()
    15  
    16  var fns = map[string]interface{}{
    17  	"quote":   regexp.QuoteMeta,
    18  	"match":   match,
    19  	"find":    find,
    20  	"replace": replace,
    21  	"split":   split,
    22  }
    23  
    24  type matchOpts struct{ Posix bool }
    25  
    26  func (*matchOpts) SetDefaultOptions() {}
    27  
    28  func match(opts matchOpts, argPattern, source string) (bool, error) {
    29  	pattern, err := makePattern(argPattern, opts.Posix, false)
    30  	if err != nil {
    31  		return false, err
    32  	}
    33  	return pattern.MatchString(source), nil
    34  }
    35  
    36  // Struct for holding options to find. Also used by split.
    37  type findOpts struct {
    38  	Posix   bool
    39  	Longest bool
    40  	Max     int
    41  }
    42  
    43  func (o *findOpts) SetDefaultOptions() { o.Max = -1 }
    44  
    45  func find(fm *eval.Frame, opts findOpts, argPattern, source string) error {
    46  	out := fm.OutputChan()
    47  
    48  	pattern, err := makePattern(argPattern, opts.Posix, opts.Longest)
    49  	if err != nil {
    50  		return err
    51  	}
    52  	matches := pattern.FindAllSubmatchIndex([]byte(source), opts.Max)
    53  
    54  	for _, match := range matches {
    55  		start, end := match[0], match[1]
    56  		groups := vector.Empty
    57  		for i := 0; i < len(match); i += 2 {
    58  			start, end := match[i], match[i+1]
    59  			text := ""
    60  			// FindAllSubmatchIndex may return negative indices to indicate
    61  			// that the pattern didn't appear in the text.
    62  			if start >= 0 && end >= 0 {
    63  				text = source[start:end]
    64  			}
    65  			groups = groups.Cons(submatchStruct{text, start, end})
    66  		}
    67  		out <- matchStruct{source[start:end], start, end, groups}
    68  	}
    69  	return nil
    70  }
    71  
    72  type replaceOpts struct {
    73  	Posix   bool
    74  	Longest bool
    75  	Literal bool
    76  }
    77  
    78  func (*replaceOpts) SetDefaultOptions() {}
    79  
    80  func replace(fm *eval.Frame, opts replaceOpts, argPattern string, argRepl interface{}, source string) (string, error) {
    81  
    82  	pattern, err := makePattern(argPattern, opts.Posix, opts.Longest)
    83  	if err != nil {
    84  		return "", err
    85  	}
    86  
    87  	if opts.Literal {
    88  		repl, ok := argRepl.(string)
    89  		if !ok {
    90  			return "", fmt.Errorf(
    91  				"replacement must be string when literal is set, got %s",
    92  				vals.Kind(argRepl))
    93  		}
    94  		return pattern.ReplaceAllLiteralString(source, repl), nil
    95  	}
    96  	switch repl := argRepl.(type) {
    97  	case string:
    98  		return pattern.ReplaceAllString(source, repl), nil
    99  	case eval.Callable:
   100  		var errReplace error
   101  		replFunc := func(s string) string {
   102  			if errReplace != nil {
   103  				return ""
   104  			}
   105  			values, err := fm.CaptureOutput(func(fm *eval.Frame) error {
   106  				return repl.Call(fm, []interface{}{s}, eval.NoOpts)
   107  			})
   108  			if err != nil {
   109  				errReplace = err
   110  				return ""
   111  			}
   112  			if len(values) != 1 {
   113  				errReplace = fmt.Errorf("replacement function must output exactly one value, got %d", len(values))
   114  				return ""
   115  			}
   116  			output, ok := values[0].(string)
   117  			if !ok {
   118  				errReplace = fmt.Errorf(
   119  					"replacement function must output one string, got %s",
   120  					vals.Kind(values[0]))
   121  				return ""
   122  			}
   123  			return output
   124  		}
   125  		return pattern.ReplaceAllStringFunc(source, replFunc), errReplace
   126  	default:
   127  		return "", fmt.Errorf(
   128  			"replacement must be string or function, got %s",
   129  			vals.Kind(argRepl))
   130  	}
   131  }
   132  
   133  func split(fm *eval.Frame, opts findOpts, argPattern, source string) error {
   134  	out := fm.OutputChan()
   135  
   136  	pattern, err := makePattern(argPattern, opts.Posix, opts.Longest)
   137  	if err != nil {
   138  		return err
   139  	}
   140  
   141  	pieces := pattern.Split(source, opts.Max)
   142  	for _, piece := range pieces {
   143  		out <- piece
   144  	}
   145  	return nil
   146  }
   147  
   148  func makePattern(p string, posix, longest bool) (*regexp.Regexp, error) {
   149  	pattern, err := compile(p, posix)
   150  	if err != nil {
   151  		return nil, err
   152  	}
   153  	if longest {
   154  		pattern.Longest()
   155  	}
   156  	return pattern, nil
   157  }
   158  
   159  func compile(pattern string, posix bool) (*regexp.Regexp, error) {
   160  	if posix {
   161  		return regexp.CompilePOSIX(pattern)
   162  	}
   163  	return regexp.Compile(pattern)
   164  }