-
Notifications
You must be signed in to change notification settings - Fork 4
/
util.go
156 lines (137 loc) · 4.16 KB
/
util.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
package common
import (
"bytes"
"encoding/json"
"fmt"
"io/ioutil"
"os"
"path/filepath"
"strings"
"github.com/fuseml/fuseml-core/pkg/util"
"github.com/pkg/errors"
"github.com/spf13/viper"
)
// KeyValueArgs is used with key-value command line arguments
type KeyValueArgs struct {
Packed []string
Unpacked map[string]string
}
// CheckErr prints a user friendly error to STDERR and exits with a non-zero
// exit code. This function is used as a wrapper for the set of steps that comprise
// the execution of a cobra command. It is the common exit point used by
// all cobra `Run` handlers. This convention, in combination with the fact that
// cobra commands only use `Run` handlers, but not `RunE` handlers (i.e. they
// don't return errors back to the cobra framework), allows for better control
// over where and how errors are handled.
func CheckErr(err error) {
if err == nil {
return
}
msg := err.Error()
if !strings.HasPrefix(msg, "error: ") {
msg = fmt.Sprintf("error: %s", msg)
}
fmt.Fprintln(os.Stderr, msg)
os.Exit(-1)
}
// Unpack converts a list of strings into a map. This helper function can be used
// to unpack command line arguments used to supplye dictionary values, e.g.:
//
// --label foo:bar --label fan: --label fin
//
// can be collected as an array of strings:
//
// ["foo:bar", "fan:", "fin"]
//
// and then unpacked with this function into a corresponding map:
//
// {"foo": "bar", "fan": "", "fin":""}
//
func (args *KeyValueArgs) Unpack() {
args.Unpacked = make(map[string]string)
for _, l := range args.Packed {
var k, v string
l = strings.TrimSpace(l)
s := strings.Split(l, ":")
if len(s) > 1 {
k = strings.TrimSpace(s[0])
v = strings.TrimSpace(strings.Join(s[1:], ":"))
} else if len(s) == 1 {
k = strings.TrimSpace(s[0])
v = ""
} else {
k = l
v = ""
}
args.Unpacked[k] = v
}
}
// LoadFileIntoVar loads the entire contents of a file into the supplied string variable
func LoadFileIntoVar(filePath string, destContent *string) error {
content, err := ioutil.ReadFile(filePath)
if err != nil {
return fmt.Errorf("cannot read file %s: %w", filePath, err)
}
*destContent = string(content)
return nil
}
// return the file handle of a config file
// if it does not exist yet, creates a new one at default location
func getCurrentOrNewConfigFile() (string, error) {
cf := viper.ConfigFileUsed()
if cf == "" {
fullname := ConfigFileName + "." + ConfigFileType
if dirname, err := os.UserHomeDir(); err == nil {
cf = filepath.Join(dirname, ConfigHomeSubdir, ConfigFuseMLSubdir, fullname)
}
if cf == "" {
return "", errors.New("Failed to acquire config directory name")
}
configDirPath := filepath.Dir(cf)
if err := os.MkdirAll(configDirPath, os.ModePerm); err != nil {
return "", err
}
fmt.Printf("FuseML configuration file created at %s\n", cf)
}
return cf, nil
}
// WriteConfigFile writes new content of the config file.
// If the file does not exist, it is created at default location
// TODO temporary solution until upstream https://github.com/spf13/viper/issues/433 is fixed
func WriteConfigFile() error {
cf, err := getCurrentOrNewConfigFile()
if err != nil {
return err
}
if err := viper.WriteConfigAs(cf); err != nil {
return err
}
return nil
}
// Writes the current content of file and also deletes given key from it
// As viper does not support this operation directly, here's a workaround
// taken from https://github.com/spf13/viper/issues/632
func DeleteKeyAndWriteConfigFile(key string) error {
cf, err := getCurrentOrNewConfigFile()
if err != nil {
return err
}
configMap := viper.AllSettings()
delete(configMap, strings.ToLower(key))
encodedConfig, _ := json.MarshalIndent(configMap, "", " ")
err = viper.ReadConfig(bytes.NewReader(encodedConfig))
if err != nil {
return err
}
if err := viper.WriteConfigAs(cf); err != nil {
return err
}
return nil
}
// ValidateEnumArgument is used to validate command line arguments that can take a limited set of values
func ValidateEnumArgument(argName, argValue string, values []string) error {
if !util.StringInSlice(argValue, values) {
return fmt.Errorf("%s must be one of: %s", argName, strings.Join(values, ", "))
}
return nil
}