go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/vpython/application/flag.go (about)

     1  // Copyright 2022 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package application
    16  
    17  import (
    18  	"flag"
    19  	"strings"
    20  
    21  	"go.chromium.org/luci/common/errors"
    22  )
    23  
    24  // boolFlag is an interface implemented by boolean flag Value instances. We use
    25  // this to determine if a flag is a boolean flag.
    26  //
    27  // This is copied from:
    28  // https://github.com/golang/go/blob/8d63408f4688ff577c25f07a1728fe131d0cae2a/src/flag/flag.go#L101
    29  type boolFlag interface {
    30  	flag.Value
    31  	IsBoolFlag() bool
    32  }
    33  
    34  func extractFlagsForSet(guardPrefix string, args []string, fs *flag.FlagSet) (fsArgs, remainder []string, err error) {
    35  	// Fast paths.
    36  	switch {
    37  	case len(args) == 0:
    38  		return
    39  	case len(args[0]) == 0 || args[0][0] != '-':
    40  		// If our first argument isn't a flag, then everything belongs to
    41  		// "remainder", no processing necessary.
    42  		remainder = args
    43  		return
    44  	}
    45  
    46  	// Scan "args" for a "--" divider. We only process candidates to the left of
    47  	// the divider.
    48  	candidates := args
    49  	for i, arg := range args {
    50  		if arg == "--" {
    51  			candidates = args[:i]
    52  			break
    53  		}
    54  	}
    55  	if len(candidates) == 0 {
    56  		remainder = args
    57  		return
    58  	}
    59  
    60  	// Make a map of all registered flags in "fs". The value will be "true" if
    61  	// the flag is a boolean flag, and "false" if it is not.
    62  	flags := make(map[string]bool)
    63  	fs.VisitAll(func(f *flag.Flag) {
    64  		bf, ok := f.Value.(boolFlag)
    65  		flags[f.Name] = ok && bf.IsBoolFlag()
    66  	})
    67  
    68  	processOne := func(args []string) (int, error) {
    69  		if len(args) == 0 {
    70  			return 0, nil
    71  		}
    72  		arg := args[0]
    73  
    74  		numMinuses := 0
    75  		if len(arg) > 0 && arg[0] == '-' {
    76  			if len(arg) > 1 && arg[1] == '-' {
    77  				numMinuses = 2
    78  			} else {
    79  				numMinuses = 1
    80  			}
    81  		}
    82  		arg = arg[numMinuses:]
    83  
    84  		if numMinuses == 0 || len(arg) == 0 {
    85  			// Not a flag, so we're done.
    86  			return 0, nil
    87  		}
    88  
    89  		single := false
    90  		eqIdx := strings.IndexRune(arg, '=')
    91  		if eqIdx >= 0 {
    92  			single = true
    93  			arg = arg[:eqIdx]
    94  		}
    95  
    96  		flagIsBool, ok := flags[arg]
    97  		if !ok {
    98  			// Unknown flag.
    99  			if strings.HasPrefix(arg, guardPrefix) {
   100  				return 0, errors.Reason("unknown flag: %s", arg).Err()
   101  			}
   102  			return 0, nil
   103  		}
   104  		if flagIsBool || single {
   105  			return 1, nil
   106  		}
   107  		return 2, nil
   108  	}
   109  
   110  	for i := 0; i < len(candidates); {
   111  		var consume int
   112  		if consume, err = processOne(candidates[i:]); err != nil {
   113  			return nil, nil, err
   114  		}
   115  
   116  		if consume == 0 {
   117  			fsArgs, remainder = args[:i], args[i:]
   118  			return
   119  		}
   120  		i += consume
   121  	}
   122  
   123  	// Got to the end with no non-"fs" flags, so everything goes to fsArgs.
   124  	fsArgs, remainder = candidates, args[len(candidates):]
   125  	return
   126  }