blob: 8815b63d69fe5e06a552da19a8a9102c51a710ed [file] [log] [blame]
// Copyright 2017 The Chromium Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
package main
import (
"strings"
"go.chromium.org/luci/common/errors"
)
// FlagSplitterDef is the definition for a compiled FlagSplitter.
type FlagSplitterDef struct {
// Solitary is the list of strings that identify as flags with no additional
// arguments.
//
// For example:
// --flag value
Solitary []string
// SolitaryAllowConjoined is the list of strings that identify as flags with
// either no additional arguments or one conjoined argument.
//
// For example:
// --flag
// --flag=value
SolitaryAllowConjoined []string
// WithArg is the list of strings that identify as flags and always have an
// additional argument.
//
// For example:
// --flag value
WithArg []string
// WithArgAllowConjoined is a special case of TwoArgs that is allowed to
// either have one additional argument or be joined to an additional argument
// via an "=".
//
// For example:
// --flag value
// --flag=value
WithArgAllowConjoined []string
}
// Compile compiles d into a FlagSplitter.
func (d FlagSplitterDef) Compile() FlagSplitter {
// Compile our total flag set into a quick lookup map.
fs := make(FlagSplitter, len(d.Solitary)+len(d.WithArg)+len(d.WithArgAllowConjoined))
fs.registerFlags(d.Solitary, flagSplitterFlag{NumArgs: 0, AllowConjoined: false})
fs.registerFlags(d.SolitaryAllowConjoined, flagSplitterFlag{NumArgs: 0, AllowConjoined: true})
fs.registerFlags(d.WithArg, flagSplitterFlag{NumArgs: 1, AllowConjoined: false})
fs.registerFlags(d.WithArgAllowConjoined, flagSplitterFlag{NumArgs: 1, AllowConjoined: true})
return fs
}
// flagSplitterFlag is a flag definition for FlagSplitter.
type flagSplitterFlag struct {
NumArgs int
AllowConjoined bool
}
// FlagSplitter contains a compiled mapping of command-line flags and their
// handling details. It is used via Parse to split a set of command-line
// argument.
type FlagSplitter map[string]flagSplitterFlag
// ParsedFlags is the output of a FlagSplitter's Split function.
type ParsedFlags struct {
// Flags is the set of parsed flags. Flags with arguments will have those
// argument as values.
Flags map[string]string
// Pos are identified positional arguments.
Pos []string
// Extra is flag data that could not be otherwise categorized. When flags
// don't match expected, we stop parsing and fill Extra with the remainder,
// since we can't otherwise know what to do with them.
Extra []string
}
func (fs FlagSplitter) registerFlags(flags []string, fsf flagSplitterFlag) {
for _, v := range flags {
fs[v] = fsf
}
}
// Split parses a set of command-line arguments into recognized flags and an
// unrecognized remainder, extra.
func (fs FlagSplitter) Split(args []string, stopAtFirstPositional bool) (pf ParsedFlags) {
for len(args) > 0 {
arg := args[0]
if !strings.HasPrefix(arg, "-") {
// Not a flag.
if stopAtFirstPositional {
pf.Pos = append(pf.Pos, args...)
return
}
// Single unknown positional argument, consume and resume parsing.
pf.Pos = append(pf.Pos, arg)
args = args[1:]
continue
}
// Hard separator.
if arg == "--" {
pf.Pos = append(pf.Pos, args[1:]...)
return
}
// If there is an "=" in the arg, then this is a conjoined flag. Store the
// conjoined part in "flagArg".
flagArg := ""
if idx := strings.IndexRune(arg, '='); idx > 0 {
arg, flagArg = arg[:idx], arg[idx+len("="):]
}
fsf, ok := fs[arg]
if !ok || flagArg != "" && !fsf.AllowConjoined {
// Either:
// - Not a known flag, so consider it (and everything following it) an
// extra argument.
// - Encountered a conjoined argument, but this type of argument is not
// allowed to be conjoined.
//
// We don't know how to handle this flag, so throw it all in "Extra".
pf.Extra = args
return
}
switch fsf.NumArgs {
case 0:
args = args[1:]
case 1:
switch {
case flagArg != "":
args = args[1:] // Conjoined
case len(args) >= 2:
flagArg = args[1]
args = args[2:]
default:
// Two-argument flag, but only one argument available. Dump the rest
// into "Extra".
pf.Extra = args
return
}
default:
panic(errors.Reason("don't know how to handle flag (%s) with more than one arg", arg).Err())
}
if pf.Flags == nil {
pf.Flags = make(map[string]string, len(args))
}
pf.Flags[arg] = flagArg
}
return
}