「 Golang 」优雅处理枚举类型

---

Open Source

https://github.com/vmware-tanzu/velero（Backup and migrate Kubernetes applications and their persistent volumes）。是用于备份与恢复集群资源的工具，支持命令行操作。其中关于命令行传入的枚举参数作了优雅校验。

Sample

https://github.com/shenxianghong/shenxianghong.github.io/tree/main/elegant-coding/enum

Structure

enum.go

package flag

import (
	"github.com/pkg/errors"
)

// Enum is a Cobra-compatible wrapper for defining a string flag that can be one of a specified set of values.
type Enum struct {
	allowedValues []string
	value         string
}

// NewEnum returns a new enum flag with the specified list of allowed values, and the specified default value if none is set.
func NewEnum(defaultValue string, allowedValues ...string) *Enum {
	return &Enum{
		allowedValues: allowedValues,
		value:         defaultValue,
	}
}

// String returns a string representation of the enum flag.
func (e *Enum) String() string {
	return e.value
}

// Set assigns the provided string to the enum receiver. It returns an error if the string is not an allowed value.
func (e *Enum) Set(s string) error {
	for _, val := range e.allowedValues {
		if val == s {
			e.value = s
			return nil
		}
	}

	return errors.Errorf("invalid value: %q", s)
}

// Type returns a string representation of the Enum type.
func (e *Enum) Type() string {
	// we don't want the help text to display anything regarding
	// the type because the usage text for the flag should capture
	// the possible options.
	return ""
}

// AllowedValues returns a slice of the flag's valid values.
func (e *Enum) AllowedValues() []string {
	return e.allowedValues
}

封装了枚举类型，包含默认值以及允许的枚举值，Enum 结构体实现了 pflag 的 Value 接口（String()，Set()，Type()），在执行 flags.StringSliceVar 等操作时，pflag 会调用 Set 函数设置值，因此 Set 可以用于校验枚举值的合法性。

format_flag.go

定义了日志格式的枚举值、默认值以及对外的工厂函数。

package main

// Format is a string representation of the desired output format for logs
type Format string

const (
	FormatText   Format = "text"
	FormatJSON   Format = "json"
	defaultValue Format = FormatText
)

// FormatFlag is a command-line flag for setting the logrus
// log format.
type FormatFlag struct {
	*Enum
	defaultValue Format
}

// NewFormatFlag constructs a new log level flag.
func NewFormatFlag() *FormatFlag {
	return &FormatFlag{
		Enum: NewEnum(
			string(defaultValue),
			string(FormatText),
			string(FormatJSON),
		),
		defaultValue: defaultValue,
	}
}

// Parse returns the flag's value as a Format.
func (f *FormatFlag) Parse() Format {
	return Format(f.String())
}

level_flags.go

原理类似 format_flag.go。

default_logger.go

package main

import (
	"github.com/sirupsen/logrus"
	"os"
)

// DefaultLogger returns a Logger with the default properties.
// The desired output format is passed as a LogFormat Enum.
func DefaultLogger(level logrus.Level, format Format) *logrus.Logger {
	logger := logrus.New()

	if format == FormatJSON {
		logger.Formatter = new(logrus.JSONFormatter)
	}

	// Make sure the output is set to stdout so log messages don't show up as errors in cloud log dashboards.
	logger.Out = os.Stdout

	logger.Level = level

	return logger
}

根据传入的枚举值，初始化一些上层使用的对象。

main.go

上层用法参考。

package main

import (
	"fmt"
	"strings"
)

func main() {
	logLevelFlag := LogLevelFlag()
	formatFlag := NewFormatFlag()

	fmt.Printf("The level at which to log. Valid values are %s.\n", strings.Join(logLevelFlag.AllowedValues(), ", "))
	fmt.Printf("The format for log output. Valid values are %s.\n", strings.Join(formatFlag.AllowedValues(), ", "))

	logLevel := logLevelFlag.Parse()
	format := FormatJSON
	
	logger := DefaultLogger(logLevel, format)
	logger.Infof("Hello World")
}