flags.go

Documentation: google.golang.org/grpc/benchmark/flags

     1  /*
     2   *
     3   * Copyright 2019 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  /*
    20  Package flags provide convenience types and routines to accept specific types
    21  of flag values on the command line.
    22  */
    23  package flags
    24  
    25  import (
    26  	"bytes"
    27  	"encoding/csv"
    28  	"flag"
    29  	"fmt"
    30  	"strconv"
    31  	"strings"
    32  	"time"
    33  )
    34  
    35  // stringFlagWithAllowedValues represents a string flag which can only take a
    36  // predefined set of values.
    37  type stringFlagWithAllowedValues struct {
    38  	val     string
    39  	allowed []string
    40  }
    41  
    42  // StringWithAllowedValues returns a flag variable of type
    43  // stringFlagWithAllowedValues configured with the provided parameters.
    44  // 'allowed` is the set of values that this flag can be set to.
    45  func StringWithAllowedValues(name, defaultVal, usage string, allowed []string) *string {
    46  	as := &stringFlagWithAllowedValues{defaultVal, allowed}
    47  	flag.CommandLine.Var(as, name, usage)
    48  	return &as.val
    49  }
    50  
    51  // String implements the flag.Value interface.
    52  func (as *stringFlagWithAllowedValues) String() string {
    53  	return as.val
    54  }
    55  
    56  // Set implements the flag.Value interface.
    57  func (as *stringFlagWithAllowedValues) Set(val string) error {
    58  	for _, a := range as.allowed {
    59  		if a == val {
    60  			as.val = val
    61  			return nil
    62  		}
    63  	}
    64  	return fmt.Errorf("want one of: %v", strings.Join(as.allowed, ", "))
    65  }
    66  
    67  type durationSliceValue []time.Duration
    68  
    69  // DurationSlice returns a flag representing a slice of time.Duration objects.
    70  func DurationSlice(name string, defaultVal []time.Duration, usage string) *[]time.Duration {
    71  	ds := make([]time.Duration, len(defaultVal))
    72  	copy(ds, defaultVal)
    73  	dsv := (*durationSliceValue)(&ds)
    74  	flag.CommandLine.Var(dsv, name, usage)
    75  	return &ds
    76  }
    77  
    78  // Set implements the flag.Value interface.
    79  func (dsv *durationSliceValue) Set(s string) error {
    80  	ds := strings.Split(s, ",")
    81  	var dd []time.Duration
    82  	for _, n := range ds {
    83  		d, err := time.ParseDuration(n)
    84  		if err != nil {
    85  			return err
    86  		}
    87  		dd = append(dd, d)
    88  	}
    89  	*dsv = durationSliceValue(dd)
    90  	return nil
    91  }
    92  
    93  // String implements the flag.Value interface.
    94  func (dsv *durationSliceValue) String() string {
    95  	var b bytes.Buffer
    96  	for i, d := range *dsv {
    97  		if i > 0 {
    98  			b.WriteRune(',')
    99  		}
   100  		b.WriteString(d.String())
   101  	}
   102  	return b.String()
   103  }
   104  
   105  type intSliceValue []int
   106  
   107  // IntSlice returns a flag representing a slice of ints.
   108  func IntSlice(name string, defaultVal []int, usage string) *[]int {
   109  	is := make([]int, len(defaultVal))
   110  	copy(is, defaultVal)
   111  	isv := (*intSliceValue)(&is)
   112  	flag.CommandLine.Var(isv, name, usage)
   113  	return &is
   114  }
   115  
   116  // Set implements the flag.Value interface.
   117  func (isv *intSliceValue) Set(s string) error {
   118  	is := strings.Split(s, ",")
   119  	var ret []int
   120  	for _, n := range is {
   121  		i, err := strconv.Atoi(n)
   122  		if err != nil {
   123  			return err
   124  		}
   125  		ret = append(ret, i)
   126  	}
   127  	*isv = intSliceValue(ret)
   128  	return nil
   129  }
   130  
   131  // String implements the flag.Value interface.
   132  func (isv *intSliceValue) String() string {
   133  	var b bytes.Buffer
   134  	for i, n := range *isv {
   135  		if i > 0 {
   136  			b.WriteRune(',')
   137  		}
   138  		b.WriteString(strconv.Itoa(n))
   139  	}
   140  	return b.String()
   141  }
   142  
   143  type stringSliceValue []string
   144  
   145  // StringSlice returns a flag representing a slice of strings.
   146  func StringSlice(name string, defaultVal []string, usage string) *[]string {
   147  	ss := make([]string, len(defaultVal))
   148  	copy(ss, defaultVal)
   149  	ssv := (*stringSliceValue)(&ss)
   150  	flag.CommandLine.Var(ssv, name, usage)
   151  	return &ss
   152  }
   153  
   154  // escapedCommaSplit splits a comma-separated list of strings in the same way
   155  // CSV files work (escaping a comma requires double-quotes).
   156  func escapedCommaSplit(str string) ([]string, error) {
   157  	r := csv.NewReader(strings.NewReader(str))
   158  	ret, err := r.Read()
   159  	if err != nil {
   160  		return nil, err
   161  	}
   162  	return ret, nil
   163  }
   164  
   165  // Set implements the flag.Value interface.
   166  func (ss *stringSliceValue) Set(str string) error {
   167  	var err error
   168  	*ss, err = escapedCommaSplit(str)
   169  	if err != nil {
   170  		return err
   171  	}
   172  	return nil
   173  }
   174  
   175  // String implements the flag.Value interface.
   176  func (ss *stringSliceValue) String() string {
   177  	return strings.Join(*ss, ",")
   178  }
   179
View as plain text