github.com/bazelbuild/remote-apis-sdks@v0.0.0-20240425170053-8a36686a6350/go/pkg/moreflag/moreflag.go (about)

     1  // Package moreflag contains definitions for some useful flag types, such as maps.
     2  package moreflag
     3  
     4  import (
     5  	"bytes"
     6  	"flag"
     7  	"fmt"
     8  	"os"
     9  	"sort"
    10  	"strings"
    11  )
    12  
    13  // Parse parses flags which are set in environment variables using the FLAG_ prefix. That is
    14  // for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
    15  // command line value of the flag takes precedence over the environment variable value.
    16  // It also calls flag.CommandLine.Parse() to parse flags sent directly as arguments, unless flag.Parse
    17  // has been previously called.
    18  func Parse() {
    19  	if !flag.Parsed() {
    20  		ParseFromEnv()
    21  		flag.CommandLine.Parse(os.Args[1:])
    22  	}
    23  }
    24  
    25  // ParseFromEnv parses flags which are set in environment variables using the FLAG_ prefix. That is
    26  // for a flag named x, x=$FLAG_x if $FLAG_x is set. If the flag is set in the command line, the
    27  // command line value of the flag takes precedence over the environment variable value.
    28  func ParseFromEnv() {
    29  	flag.VisitAll(func(f *flag.Flag) {
    30  		v, ok := os.LookupEnv("FLAG_" + f.Name)
    31  		if ok {
    32  			flag.Set(f.Name, v)
    33  		}
    34  	})
    35  }
    36  
    37  // StringMapValue is a command line flag that interprets a string in the format key1=value1,key2=value2
    38  // as a map.
    39  type StringMapValue map[string]string
    40  
    41  // String retrieves the flag's map in the format key1=value1,key2=value, sorted by keys.
    42  func (m *StringMapValue) String() string {
    43  	// Construct the output in key sorted order
    44  	keys := make([]string, 0, len(*m))
    45  	for key := range *m {
    46  		keys = append(keys, key)
    47  	}
    48  	sort.Strings(keys)
    49  	var b bytes.Buffer
    50  	for i, key := range keys {
    51  		if i > 0 {
    52  			b.WriteRune(',')
    53  		}
    54  		b.WriteString(key)
    55  		b.WriteRune('=')
    56  		b.WriteString((*m)[key])
    57  	}
    58  	return b.String()
    59  }
    60  
    61  // Set updates the map with key and value pair(s) in the format key1=value1,key2=value2.
    62  func (m *StringMapValue) Set(s string) error {
    63  	*m = make(map[string]string)
    64  	pairs := strings.Split(s, ",")
    65  	for _, p := range pairs {
    66  		if p == "" {
    67  			continue
    68  		}
    69  		pair := strings.Split(p, "=")
    70  		if len(pair) != 2 {
    71  			return fmt.Errorf("wrong format for key-value pair: %v", p)
    72  		}
    73  		if pair[0] == "" {
    74  			return fmt.Errorf("key not provided")
    75  		}
    76  		if _, ok := (*m)[pair[0]]; ok {
    77  			return fmt.Errorf("key %v already defined in list of key-value pairs %v", pair[0], s)
    78  		}
    79  		(*m)[pair[0]] = pair[1]
    80  	}
    81  	return nil
    82  }
    83  
    84  // Get returns the flag value as a map of strings.
    85  func (m *StringMapValue) Get() interface{} {
    86  	return map[string]string(*m)
    87  }
    88  
    89  // StringListValue is a command line flag that interprets a string as a list of comma-separated values.
    90  type StringListValue []string
    91  
    92  // String returns the list value.
    93  func (m *StringListValue) String() string {
    94  	return strings.Join(*m, ",")
    95  }
    96  
    97  // Set for StringListValue accepts one list of comma-separated values.
    98  func (m *StringListValue) Set(s string) error {
    99  	splitFn := func(c rune) bool {
   100  		return c == ','
   101  	}
   102  	*m = StringListValue(strings.FieldsFunc(s, splitFn))
   103  	return nil
   104  }
   105  
   106  // Get returns the flag value as a list of strings.
   107  func (m *StringListValue) Get() interface{} {
   108  	return []string(*m)
   109  }