github.com/cockroachdb/cockroach@v20.2.0-alpha.1+incompatible/pkg/util/flagutil/flagutil.go (about)

     1  // Copyright 2018 The Cockroach Authors.
     2  //
     3  // Use of this software is governed by the Business Source License
     4  // included in the file licenses/BSL.txt.
     5  //
     6  // As of the Change Date specified in that file, in accordance with
     7  // the Business Source License, use of this software will be governed
     8  // by the Apache License, Version 2.0, included in the file
     9  // licenses/APL.txt.
    10  
    11  // Package flagutil facilitates creation of rich flag types.
    12  package flagutil
    13  
    14  import (
    15  	"fmt"
    16  	"regexp"
    17  	"time"
    18  
    19  	"github.com/cockroachdb/cockroach/pkg/util/log"
    20  	"github.com/cockroachdb/cockroach/pkg/util/timeutil"
    21  	"github.com/spf13/pflag"
    22  )
    23  
    24  // TimeFormats are the time formats which a time flag accepts for parsing time
    25  // as described in time.Parse.
    26  var TimeFormats = []string{
    27  	log.MessageTimeFormat,
    28  	log.FileTimeFormat,
    29  	time.RFC3339Nano,
    30  	time.RFC3339,
    31  	time.Kitchen,
    32  	"15:04:00.999999Z",
    33  	"15:04:00.999999",
    34  	"15:04Z",
    35  	"15:04",
    36  }
    37  
    38  // Time returns a value which can be used with pflag.Var to create flags for
    39  // time variables. The flag attempts to parse its in a variety of formats.
    40  // The first is as a duration using time.ParseDuration and subtracting the
    41  // result from Now() and then by using the values in Formats for parsing using
    42  // time.Parse. An empty flag sets t to the zero value.
    43  func Time(t *time.Time) pflag.Value {
    44  	return (*timeFlag)(t)
    45  }
    46  
    47  // Regexp returns a value which can be used with pflag.Var to create flags
    48  // for regexp variables. The flag attempts to compile its input to a regular
    49  // expression and assigns the result to *r.
    50  // If the flag is empty, r is set to nil.
    51  func Regexp(r **regexp.Regexp) pflag.Value {
    52  	return re{re: r}
    53  }
    54  
    55  // timeFlag implements pflag.Value and enables easy creation of time flags.
    56  type timeFlag time.Time
    57  
    58  func (f *timeFlag) String() string {
    59  	if (*time.Time)(f).IsZero() {
    60  		return ""
    61  	}
    62  	return (*time.Time)(f).Format(log.MessageTimeFormat)
    63  }
    64  
    65  func (f *timeFlag) Type() string { return "time" }
    66  
    67  func (f *timeFlag) Set(s string) error {
    68  	if s == "" {
    69  		*f = (timeFlag)(time.Time{})
    70  		return nil
    71  	}
    72  	for _, p := range parsers {
    73  		if t, err := p.parseTime(s); err == nil {
    74  			*f = timeFlag(t.UTC())
    75  			return nil
    76  		}
    77  	}
    78  	return fmt.Errorf("failed to parse %q as time", s)
    79  }
    80  
    81  var parsers = func() (parsers []interface {
    82  	parseTime(s string) (time.Time, error)
    83  }) {
    84  	parsers = append(parsers, durationParser(time.ParseDuration))
    85  	for _, p := range TimeFormats {
    86  		parsers = append(parsers, formatParser(p))
    87  	}
    88  	return parsers
    89  }()
    90  
    91  type formatParser string
    92  
    93  func (f formatParser) parseTime(s string) (time.Time, error) {
    94  	return time.Parse(string(f), s)
    95  }
    96  
    97  type durationParser func(string) (time.Duration, error)
    98  
    99  func (f durationParser) parseTime(s string) (time.Time, error) {
   100  	d, err := f(s)
   101  	if err != nil {
   102  		return time.Time{}, err
   103  	}
   104  	return timeutil.Now().Add(-1 * d), nil
   105  }
   106  
   107  type re struct {
   108  	re **regexp.Regexp
   109  }
   110  
   111  func (r re) String() string {
   112  	if *r.re == nil {
   113  		return ""
   114  	}
   115  	return (*r.re).String()
   116  }
   117  
   118  func (r re) Set(s string) error {
   119  	if s == "" {
   120  		*r.re = nil
   121  		return nil
   122  	}
   123  	re, err := regexp.Compile(s)
   124  	if err != nil {
   125  		return err
   126  	}
   127  	*r.re = re
   128  	return nil
   129  }
   130  
   131  func (r re) Type() string { return "regexp" }