From 0b4d636ec2fcbc394c119e3c2804d52aa6dafdea Mon Sep 17 00:00:00 2001 From: Adphi Date: Tue, 17 Oct 2023 15:30:32 +0200 Subject: [PATCH] add cobra command utilities and log formatter Signed-off-by: Adphi --- cli/clifmt/formatter.go | 84 ++++++++++++ cli/command.go | 297 ++++++++++++++++++++++++++++++++++++++++ cli/command_test.go | 49 +++++++ 3 files changed, 430 insertions(+) create mode 100644 cli/clifmt/formatter.go create mode 100644 cli/command.go create mode 100644 cli/command_test.go diff --git a/cli/clifmt/formatter.go b/cli/clifmt/formatter.go new file mode 100644 index 0000000..65cf66c --- /dev/null +++ b/cli/clifmt/formatter.go @@ -0,0 +1,84 @@ +// Copyright 2023 Linka Cloud All rights reserved. +// +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package clifmt + +import ( + "bytes" + "strings" + "time" + + "github.com/fatih/color" + "github.com/sirupsen/logrus" +) + +type TimeFormat string + +const ( + NoneTimeFormat TimeFormat = "none" + FullTimeFormat TimeFormat = "full" + RelativeTimeFormat TimeFormat = "relative" +) + +const ( + red = 31 + yellow = 33 + blue = 36 + white = 39 + gray = 90 +) + +func New(f TimeFormat) logrus.Formatter { + return &clifmt{start: time.Now(), format: f} +} + +type clifmt struct { + start time.Time + format TimeFormat +} + +func (f *clifmt) Format(entry *logrus.Entry) ([]byte, error) { + var b bytes.Buffer + var c *color.Color + switch entry.Level { + case logrus.DebugLevel, logrus.TraceLevel: + c = color.New(gray) + case logrus.WarnLevel: + c = color.New(yellow) + case logrus.ErrorLevel, logrus.FatalLevel, logrus.PanicLevel: + c = color.New(red) + default: + c = color.New(white) + } + msg := entry.Message + if len(entry.Message) > 0 && entry.Level < logrus.DebugLevel { + msg = strings.ToTitle(string(msg[0])) + msg[1:] + } + + var err error + switch f.format { + case FullTimeFormat: + _, err = c.Fprintf(&b, "[%s] %s\n", entry.Time.Format("2006-01-02 15:04:05"), entry.Message) + case RelativeTimeFormat: + _, err = c.Fprintf(&b, "[%5v] %s\n", entry.Time.Sub(f.start).Truncate(time.Second).String(), msg) + case NoneTimeFormat: + fallthrough + default: + _, err = c.Fprintln(&b, msg) + } + if err != nil { + return nil, err + } + return b.Bytes(), nil +} diff --git a/cli/command.go b/cli/command.go new file mode 100644 index 0000000..dcf5136 --- /dev/null +++ b/cli/command.go @@ -0,0 +1,297 @@ +// Package cli is adapted from https://github.com/rancher/wrangler-cli +package cli + +import ( + "fmt" + "os" + "reflect" + "regexp" + "strconv" + "strings" + "unsafe" + + "github.com/sirupsen/logrus" + "github.com/spf13/cobra" + "github.com/spf13/pflag" + + "go.linka.cloud/grpc-toolkit/signals" +) + +var ( + caseRegexp = regexp.MustCompile("([a-z])([A-Z])") +) + +type PersistentPreRunnable interface { + PersistentPre(cmd *cobra.Command, args []string) error +} + +type PreRunnable interface { + Pre(cmd *cobra.Command, args []string) error +} + +type Runnable interface { + Run(cmd *cobra.Command, args []string) error +} + +type customizer interface { + Customize(cmd *cobra.Command) +} + +type fieldInfo struct { + FieldType reflect.StructField + FieldValue reflect.Value +} + +func fields(obj interface{}) []fieldInfo { + ptrValue := reflect.ValueOf(obj) + objValue := ptrValue.Elem() + + var result []fieldInfo + + for i := 0; i < objValue.NumField(); i++ { + fieldType := objValue.Type().Field(i) + if !fieldType.IsExported() { + continue + } + if fieldType.Anonymous && fieldType.Type.Kind() == reflect.Struct { + result = append(result, fields(objValue.Field(i).Addr().Interface())...) + } else if !fieldType.Anonymous { + result = append(result, fieldInfo{ + FieldValue: objValue.Field(i), + FieldType: objValue.Type().Field(i), + }) + } + } + + return result +} + +func Name(obj interface{}) string { + ptrValue := reflect.ValueOf(obj) + objValue := ptrValue.Elem() + commandName := strings.Replace(objValue.Type().Name(), "Command", "", 1) + commandName, _ = name(commandName, "", "") + return commandName +} + +func Main(cmd *cobra.Command) { + ctx := signals.SetupSignalHandler() + if err := cmd.ExecuteContext(ctx); err != nil { + logrus.Fatal(err) + } +} + +func makeEnvVar[T comparable](to []func(), name string, vars []string, defValue T, flags *pflag.FlagSet, fn func(flag string) (T, error)) []func() { + for _, v := range vars { + to = append(to, func() { + v := os.Getenv(v) + if v == "" { + return + } + fv, err := fn(name) + if err == nil && fv == defValue { + flags.Set(name, v) + } + }) + } + return to +} + +// Command populates a obj.Command() object by extracting args from struct tags of the +// Runnable obj passed. Also the Run method is assigned to the RunE of the cli. +// name = Override the struct field with +func Command(obj Runnable, c *cobra.Command) *cobra.Command { + var ( + envs []func() + arrays = map[string]reflect.Value{} + slices = map[string]reflect.Value{} + maps = map[string]reflect.Value{} + ptrValue = reflect.ValueOf(obj) + objValue = ptrValue.Elem() + ) + + if c.Use == "" { + c.Use = Name(obj) + } + + for _, info := range fields(obj) { + fieldType := info.FieldType + v := info.FieldValue + + name, alias := name(fieldType.Name, fieldType.Tag.Get("name"), fieldType.Tag.Get("short")) + usage := fieldType.Tag.Get("usage") + envVars := strings.Split(fieldType.Tag.Get("env"), ",") + defValue := fieldType.Tag.Get("default") + if len(envVars) == 1 && envVars[0] == "" { + envVars = nil + } + for _, v := range envVars { + if v == "" { + continue + } + usage += fmt.Sprintf("%s [$%s]", usage, v) + } + defInt, err := strconv.Atoi(defValue) + if err != nil { + defInt = 0 + } + defValueLower := strings.ToLower(defValue) + defBool := defValueLower == "true" || defValueLower == "1" || defValueLower == "yes" || defValueLower == "y" + + flags := c.PersistentFlags() + switch fieldType.Type.Kind() { + case reflect.Int: + flags.IntVarP((*int)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defInt, usage) + envs = append(envs, makeEnvVar(envs, name, envVars, defInt, flags, flags.GetInt)...) + case reflect.String: + flags.StringVarP((*string)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defValue, usage) + envs = append(envs, makeEnvVar(envs, name, envVars, defValue, flags, flags.GetString)...) + case reflect.Slice: + // env is not supported for slices + switch fieldType.Tag.Get("split") { + case "false": + arrays[name] = v + flags.StringArrayP(name, alias, nil, usage) + default: + slices[name] = v + flags.StringSliceP(name, alias, nil, usage) + } + case reflect.Map: + maps[name] = v + flags.StringSliceP(name, alias, nil, usage) + case reflect.Bool: + flags.BoolVarP((*bool)(unsafe.Pointer(v.Addr().Pointer())), name, alias, defBool, usage) + envs = append(envs, makeEnvVar(envs, name, envVars, defBool, flags, flags.GetBool)...) + default: + panic("Unknown kind on field " + fieldType.Name + " on " + objValue.Type().Name()) + } + + if len(envVars) == 0 { + continue + } + } + + if p, ok := obj.(PersistentPreRunnable); ok { + c.PersistentPreRunE = p.PersistentPre + } + + if p, ok := obj.(PreRunnable); ok { + c.PreRunE = p.Pre + } + + c.RunE = obj.Run + c.PersistentPreRunE = bind(c.PersistentPreRunE, arrays, slices, maps, envs) + c.PreRunE = bind(c.PreRunE, arrays, slices, maps, envs) + c.RunE = bind(c.RunE, arrays, slices, maps, envs) + + cust, ok := obj.(customizer) + if ok { + cust.Customize(c) + } + + return c +} + +func assignMaps(app *cobra.Command, maps map[string]reflect.Value) error { + for k, v := range maps { + k = contextKey(k) + s, err := app.Flags().GetStringSlice(k) + if err != nil { + return err + } + if s != nil { + values := map[string]string{} + for _, part := range s { + parts := strings.SplitN(part, "=", 2) + if len(parts) == 1 { + values[parts[0]] = "" + } else { + values[parts[0]] = parts[1] + } + } + v.Set(reflect.ValueOf(values)) + } + } + return nil +} + +func assignSlices(app *cobra.Command, slices map[string]reflect.Value) error { + for k, v := range slices { + k = contextKey(k) + s, err := app.Flags().GetStringSlice(k) + if err != nil { + return err + } + if s != nil { + v.Set(reflect.ValueOf(s[:])) + } + } + return nil +} + +func assignArrays(app *cobra.Command, arrays map[string]reflect.Value) error { + for k, v := range arrays { + k = contextKey(k) + s, err := app.Flags().GetStringArray(k) + if err != nil { + return err + } + if s != nil { + v.Set(reflect.ValueOf(s[:])) + } + } + return nil +} + +func contextKey(name string) string { + parts := strings.Split(name, ",") + return parts[len(parts)-1] +} + +func name(name, setName, short string) (string, string) { + if setName != "" { + return setName, short + } + parts := strings.Split(name, "_") + i := len(parts) - 1 + name = caseRegexp.ReplaceAllString(parts[i], "$1-$2") + name = strings.ToLower(name) + result := append([]string{name}, parts[0:i]...) + for i := 0; i < len(result); i++ { + result[i] = strings.ToLower(result[i]) + } + if short == "" && len(result) > 1 { + short = result[1] + } + return result[0], short +} + +func bind(next func(*cobra.Command, []string) error, + arrays map[string]reflect.Value, + slices map[string]reflect.Value, + maps map[string]reflect.Value, + envs []func()) func(*cobra.Command, []string) error { + if next == nil { + return nil + } + return func(cmd *cobra.Command, args []string) error { + for _, envCallback := range envs { + envCallback() + } + if err := assignArrays(cmd, arrays); err != nil { + return err + } + if err := assignSlices(cmd, slices); err != nil { + return err + } + if err := assignMaps(cmd, maps); err != nil { + return err + } + + if next != nil { + return next(cmd, args) + } + + return nil + } +} diff --git a/cli/command_test.go b/cli/command_test.go new file mode 100644 index 0000000..1b3c220 --- /dev/null +++ b/cli/command_test.go @@ -0,0 +1,49 @@ +package cli + +import ( + "os" + "testing" + + "github.com/spf13/cobra" + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +type testCmd struct { + String string `name:"string" short:"s" usage:"string flag" env:"STRING" default:"string"` + Int int `name:"int" short:"i" usage:"int flag" env:"INT" default:"1"` + Bool bool `name:"bool" short:"b" usage:"bool flag" env:"BOOL" default:"true"` + StringSlice []string `name:"string-slice" short:"S" usage:"string slice flag"` +} + +func (c *testCmd) Run(cmd *cobra.Command, args []string) error { + return nil +} + +func TestCommand(t *testing.T) { + var c testCmd + cmd := Command(&c, &cobra.Command{ + Short: "test", + }) + require.NoError(t, cmd.Execute()) + assert.Equal(t, "string", c.String) + assert.Equal(t, 1, c.Int) + assert.Equal(t, true, c.Bool) + assert.Equal(t, []string{}, c.StringSlice) +} + +func TestCommandEnv(t *testing.T) { + require.NoError(t, os.Setenv("STRING", "env-string")) + require.NoError(t, os.Setenv("INT", "2")) + require.NoError(t, os.Setenv("BOOL", "false")) + require.NoError(t, os.Setenv("STRING_SLICE", "env-string1,env-string2")) + var c testCmd + cmd := Command(&c, &cobra.Command{ + Short: "test", + }) + require.NoError(t, cmd.Execute()) + assert.Equal(t, "env-string", c.String) + assert.Equal(t, 2, c.Int) + assert.Equal(t, false, c.Bool) + assert.Equal(t, []string{}, c.StringSlice) +}