Golang FlagSet.Visit方法代码示例

func loadConfig(nm string, flagSet *flag.FlagSet, isOverride bool) error {
	f, e := os.Open(nm)
	if nil != e {
		return fmt.Errorf("load config '%s' failed, %v", nm, e)
	}

	var res map[string]interface{}
	e = json.NewDecoder(f).Decode(&res)
	if nil != e {
		return fmt.Errorf("load config '%s' failed, %v", nm, e)
	}

	actual := map[string]string{}
	fn := func(f *flag.Flag) {
		actual[f.Name] = f.Name
	}
	if nil == flagSet {
		flag.Visit(fn)
	} else {
		flagSet.Visit(fn)
	}

	e = assignFlagSet("", res, flagSet, actual, isOverride)
	if nil != e {
		return fmt.Errorf("load config '%s' failed, %v", nm, e)
	}
	return nil
}

// SetFlagsFromYaml goes through all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// the YAML config. It will use the key REPLACE(UPPERCASE(flagname), '-', '_')
func SetFlagsFromYaml(fs *flag.FlagSet, rawYaml []byte) (err error) {
	conf := make(map[string]string)
	if err = yaml.Unmarshal(rawYaml, conf); err != nil {
		return
	}
	alreadySet := map[string]struct{}{}
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = struct{}{}
	})

	errs := make([]error, 0)
	fs.VisitAll(func(f *flag.Flag) {
		if f.Name == "" {
			return
		}
		if _, ok := alreadySet[f.Name]; ok {
			return
		}
		tag := strings.Replace(strings.ToUpper(f.Name), "-", "_", -1)
		val, ok := conf[tag]
		if !ok {
			return
		}
		if serr := fs.Set(f.Name, val); serr != nil {
			errs = append(errs, fmt.Errorf("invalid value %q for %s: %v", val, tag, serr))
		}
	})
	if len(errs) != 0 {
		err = ErrorSlice(errs)
	}
	return
}

func normalizeFlags(flags []Flag, set *flag.FlagSet) error {
	visited := make(map[string]bool)
	set.Visit(func(f *flag.Flag) {
		visited[f.Name] = true
	})
	for _, f := range flags {
		parts := strings.Split(f.getName(), ",")
		if len(parts) == 1 {
			continue
		}
		var ff *flag.Flag
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if visited[name] {
				if ff != nil {
					return errors.New("Cannot use two forms of the same flag: " + name + " " + ff.Name)
				}
				ff = set.Lookup(name)
			}
		}
		if ff == nil {
			continue
		}
		for _, name := range parts {
			name = strings.Trim(name, " ")
			if !visited[name] {
				copyFlag(name, ff, set)
			}
		}
	}
	return nil
}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the prefix "PREFIX_", and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func SetFlagsFromEnv(fs *flag.FlagSet, prefix string) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})

	if prefix == "" {
		prefix = "_"
	} else {
		prefix = strings.ToUpper(strings.Replace(prefix, "-", "_", -1))
	}

	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := prefix + strings.ToUpper(strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				if serr := fs.Set(f.Name, val); serr != nil {
					err = fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
				}
			}
		}
	})
	return err
}

// ParseSet parses the given flagset. The specified prefix will be applied to
// the environment variable names.
func ParseSet(prefix string, set *flag.FlagSet) error {
	var explicit []*flag.Flag
	var all []*flag.Flag
	set.Visit(func(f *flag.Flag) {
		explicit = append(explicit, f)
	})

	var err error
	set.VisitAll(func(f *flag.Flag) {
		if err != nil {
			return
		}
		all = append(all, f)
		if !contains(explicit, f) {
			name := strings.Replace(f.Name, ".", "_", -1)
			name = strings.Replace(name, "-", "_", -1)
			if prefix != "" {
				name = prefix + name
			}
			name = strings.ToUpper(name)
			val := os.Getenv(name)
			if val != "" {
				if ferr := f.Value.Set(val); ferr != nil {
					err = fmt.Errorf("failed to set flag %q with value %q", f.Name, val)
				}
			}
		}
	})
	return err
}

func IsSet(fs *flag.FlagSet, name string) bool {
	set := false
	fs.Visit(func(f *flag.Flag) {
		if f.Name == name {
			set = true
		}
	})
	return set
}

// IsFlagSet returns true if the specified flag has been set on
// the command line.
func IsFlagSet(fs *flag.FlagSet, name string) bool {
	found := false
	fs.Visit(func(f *flag.Flag) {
		if f.Name == name {
			found = true
		}
	})
	return found
}

// FlagVisited determines whether the named flag has been visited in the FlagSet
// given. This is helpful if you want to have a flag that triggers an action
// when given, but is not a boolean flag.
func FlagVisited(name string, flags flag.FlagSet) (didFind bool) {
	toFind := flags.Lookup(name)

	flags.Visit(func(f *flag.Flag) {
		if f == toFind {
			didFind = true
		}
	})

	return
}

func loadActualFlags(flagSet *flag.FlagSet) map[string]string {
	actual := map[string]string{}
	fn := func(f *flag.Flag) {
		actual[f.Name] = f.Name
	}
	if nil == flagSet {
		flag.Visit(fn)
	} else {
		flagSet.Visit(fn)
	}
	return actual
}

// EnabledActionsFromFlags returns value as an array of string
// that is assigned to a command line flag name `enabledActionsFlagName`.
func EnabledActionsFromFlags(fs *flag.FlagSet, enabledActionsFlagName string) ([]string, error) {
	visited := make(map[string]struct{})
	fs.Visit(func(f *flag.Flag) {
		visited[f.Name] = struct{}{}
	})

	_, enabledActionsFlagIsSet := visited[enabledActionsFlagName]

	if enabledActionsFlagIsSet {
		enabledActions := *fs.Lookup(enabledActionsFlagName).Value.(*ActionNames)
		return enabledActions, nil
	}

	return []string{}, nil
}

// flagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix, and any dashes are replaced by
// underscores - for example: some-flag => PREFIX_SOME_FLAG
func flagsFromEnv(prefix string, fs *flag.FlagSet) {
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})
	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := strings.ToUpper(prefix + "_" + strings.Replace(f.Name, "-", "_", -1))
			val := os.Getenv(key)
			if val != "" {
				fs.Set(f.Name, val)
			}
		}
	})
}

// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the given prefix  and any dashes are replaced by
// underscores - for example: some-flag => ETCD_SOME_FLAG
func SetFlagsFromEnv(prefix string, fs *flag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[flagToEnv(prefix, f.Name)] = true
	})
	usedEnvKey := make(map[string]bool)
	fs.VisitAll(func(f *flag.Flag) {
		err = setFlagFromEnv(fs, prefix, f.Name, usedEnvKey, alreadySet, true)
	})

	verifyEnv(prefix, usedEnvKey, alreadySet)

	return err
}

func mustNotHave(fset *flag.FlagSet, keys ...string) error {
	for _, key := range keys {
		found := false
		fset.Visit(
			func(f *flag.Flag) {
				if f.Name == key {
					found = true
				}
			})
		if found {
			return fmt.Errorf("Invalid flags. Flag '%s' cannot appear for this operation", key)
		}
	}
	return nil
}

func mustHave(fset *flag.FlagSet, keys ...string) error {
	for _, key := range keys {
		found := false
		fset.Visit(
			func(f *flag.Flag) {
				if f.Name == key {
					found = true
				}
			})
		if !found {
			flag := fset.Lookup(key)
			if flag == nil || flag.DefValue == "" {
				return fmt.Errorf("Invalid flags. Flag '%s' is required for this operation", key)
			}
		}
	}
	return nil
}

// from github.com/coreos/etcd
//
// SetFlagsFromEnv parses all registered flags in the given flagset,
// and if they are not already set it attempts to set their values from
// environment variables. Environment variables take the name of the flag but
// are UPPERCASE, have the prefix `s`, and any dashes are replaced by
// underscores - for example: some-flag => ETCD_SOME_FLAG
func AddEnvironmentToFlags(s string, fs *flag.FlagSet) error {
	var err error
	alreadySet := make(map[string]bool)
	fs.Visit(func(f *flag.Flag) {
		alreadySet[f.Name] = true
	})
	fs.VisitAll(func(f *flag.Flag) {
		if !alreadySet[f.Name] {
			key := s + "_" + ToEnvironmentKey(f.Name)
			val := os.Getenv(key)
			if val != "" {
				if serr := fs.Set(f.Name, val); serr != nil {
					err = fmt.Errorf("invalid value %q for %s: %v", val, key, serr)
				}
			}
		}
	})
	return err
}