/
envvar.go
353 lines (319 loc) · 10.8 KB
/
envvar.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
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
// Copyright 2015 The Vanadium 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 envvar implements utilities for processing environment variables.
// There are three representations of environment variables:
// 1. []"key=value" # hard to get and set, used by standard Go packages
// 2. map[key]value # simple to get and set, nicest syntax
// 3. *envvar.Vars # simple to get and set, also tracks deltas
//
// The slice form (1) is used by standard Go packages, presumably since it's
// similar to the underlying OS representation. The map form (2) is convenient
// to use, and has native Go map syntax. The Vars form (3) is also convenient
// to use, and tracks deltas when mutations are performed.
//
// This package provides utilities to easily use and convert between the three
// representations.
//
// Empty keys are invalid and silently skipped in operations over all
// representations.
package envvar
import (
"os"
"sort"
"strings"
)
// MergeMaps merges together maps, and returns a new map with the merged result.
// If the same key appears in more than one input map, the last one "wins"; the
// value is set based on the last map containing that key.
//
// As a result of its semantics, MergeMaps called with a single map returns a
// copy of the map, with empty keys dropped.
func MergeMaps(maps ...map[string]string) map[string]string {
merged := make(map[string]string)
for _, m := range maps {
for key, value := range m {
if key != "" {
merged[key] = value
}
}
}
return merged
}
// CopyMap returns a copy of from, with empty keys dropped.
func CopyMap(from map[string]string) map[string]string {
return MergeMaps(from)
}
// MergeSlices merges together slices, and returns a new slice with the merged
// result. If the same key appears more than once in a single input slice, or
// in more than one input slice, the last one "wins"; the value is set based on
// the last slice element in the last slice containing that key.
//
// As a result of its semantics, MergeSlices called with a single slice returns
// a copy of the slice, with empty keys dropped.
func MergeSlices(slices ...[]string) []string {
merged := make(map[string]string)
for _, slice := range slices {
for _, kv := range slice {
if key, value := SplitKeyValue(kv); key != "" {
merged[key] = value
}
}
}
return MapToSlice(merged)
}
// CopySlice returns a copy of from, with empty keys dropped, and ordered by
// key. If the same key appears more than once the last one "wins"; the value
// is set based on the last slice element containing that key.
func CopySlice(from []string) []string {
return MergeSlices(from)
}
// MapToSlice converts from the map to the slice representation. The returned
// slice is in sorted order.
func MapToSlice(from map[string]string) []string {
to := make([]string, 0, len(from))
for key, value := range from {
if key != "" {
to = append(to, JoinKeyValue(key, value))
}
}
SortByKey(to)
return to
}
// SliceToMap converts from the slice to the map representation. If the same
// key appears more than once, the last one "wins"; the value is set based on
// the last slice element containing that key.
func SliceToMap(from []string) map[string]string {
to := make(map[string]string, len(from))
for _, kv := range from {
if key, value := SplitKeyValue(kv); key != "" {
to[key] = value
}
}
return to
}
// SplitKeyValue splits kv into its key and value components. The format of kv
// is "key=value"; the split is performed on the first '=' character.
func SplitKeyValue(kv string) (string, string) {
split := strings.SplitN(kv, "=", 2)
if len(split) == 2 {
return split[0], split[1]
}
return split[0], ""
}
// JoinKeyValue joins key and value into a single string "key=value".
func JoinKeyValue(key, value string) string {
return key + "=" + value
}
// SplitTokens is like strings.Split(value, separator), but also filters out
// empty tokens. Thus SplitTokens("", ":") returns a nil slice, unlike
// strings.SplitTokens which returns a slice with a single empty string.
func SplitTokens(value, separator string) []string {
var tokens []string
for _, token := range strings.Split(value, separator) {
if token != "" {
tokens = append(tokens, token)
}
}
return tokens
}
// JoinTokens is like strings.Join(tokens, separator), but also filters out
// empty tokens.
func JoinTokens(tokens []string, separator string) string {
var value string
for _, token := range tokens {
if token == "" {
continue
}
if value != "" {
value += separator
}
value += token
}
return value
}
// UniqueTokens returns a new slice containing tokens that are not empty or
// duplicated, and in the same relative order as the original slice.
func UniqueTokens(tokens []string) []string {
var unique []string
seen := make(map[string]bool)
for _, token := range tokens {
if token == "" || seen[token] {
continue
}
seen[token] = true
unique = append(unique, token)
}
return unique
}
// FilterToken returns a new slice containing tokens that are not empty or match
// the target, and in the same relative order as the original slice.
func FilterToken(tokens []string, target string) []string {
var filtered []string
for _, token := range tokens {
if token == "" || token == target {
continue
}
filtered = append(filtered, token)
}
return filtered
}
// PrependUniqueToken prepends token to value, which is separated by separator,
// removing all empty and duplicate tokens. Returns a string where token only
// occurs once, and is first.
func PrependUniqueToken(value, separator, token string) string {
result := SplitTokens(value, separator)
result = append([]string{token}, result...)
return JoinTokens(UniqueTokens(result), separator)
}
// AppendUniqueToken appends token to value, which is separated by separator,
// and removes all empty and duplicate tokens. Returns a string where token
// only occurs once, and is last.
func AppendUniqueToken(value, separator, token string) string {
result := SplitTokens(value, separator)
result = FilterToken(result, token)
result = append(result, token)
return JoinTokens(UniqueTokens(result), separator)
}
// SortByKey sorts vars into ascending key order, where vars is expected to be
// in the []"key=value" slice representation.
func SortByKey(vars []string) {
sort.Sort(keySorter(vars))
}
type keySorter []string
func (s keySorter) Len() int { return len(s) }
func (s keySorter) Swap(i, j int) { s[i], s[j] = s[j], s[i] }
func (s keySorter) Less(i, j int) bool {
ikey, _ := SplitKeyValue(s[i])
jkey, _ := SplitKeyValue(s[j])
return ikey < jkey
}
// Vars is a mutable set of environment variables that tracks deltas.
//
// Vars are initialized with a base environment, and may be mutated with calls
// to Set and SetTokens. The resulting environment is retrieved with calls to
// ToMap and ToSlice.
//
// Mutations are tracked separately from the base environment; call Deltas to
// retrieve only the environment variables that have been changed.
//
// The zero Vars has an empty environment, and supports all methods.
type Vars struct {
base map[string]string
deltaInsert map[string]string
deltaRemove map[string]bool
}
// VarsFromMap returns a new Vars initialized from the given base map.
func VarsFromMap(base map[string]string) *Vars { return &Vars{base: MergeMaps(base)} }
// VarsFromSlice returns a new Vars initialized from the given base slice.
func VarsFromSlice(base []string) *Vars { return &Vars{base: SliceToMap(base)} }
// VarsFromOS returns a new Vars initialized from os.Environ.
func VarsFromOS() *Vars { return VarsFromSlice(os.Environ()) }
// Contains returns true iff the key exists in the current set of variables.
func (x *Vars) Contains(key string) bool {
if x.deltaRemove[key] {
return false
}
if _, ok := x.deltaInsert[key]; ok {
return true
}
_, ok := x.base[key]
return ok
}
// Get returns the value associated with key. Returns "" if the key doesn't
// exist, or if the key has an empty value. Use Contains to test for existence.
func (x *Vars) Get(key string) string {
if x.deltaRemove[key] {
return ""
}
if value, ok := x.deltaInsert[key]; ok {
return value
}
return x.base[key]
}
// GetTokens is a convenience that calls SplitTokens(x.Get(key), separator).
func (x *Vars) GetTokens(key, separator string) []string {
return SplitTokens(x.Get(key), separator)
}
// Set assigns key to the given value.
func (x *Vars) Set(key, value string) {
if key != "" {
if x.deltaInsert == nil {
x.deltaInsert = make(map[string]string)
}
x.deltaInsert[key] = value
delete(x.deltaRemove, key)
}
}
// SetTokens is a convenience that calls x.Set(key, JoinTokens(tokens, separator)).
func (x *Vars) SetTokens(key string, tokens []string, separator string) {
x.Set(key, JoinTokens(tokens, separator))
}
// Delete removes the given keys. Subsequent calls to Contains on each key
// will return false.
func (x *Vars) Delete(keys ...string) {
for _, key := range keys {
if key != "" {
if x.deltaRemove == nil {
x.deltaRemove = make(map[string]bool)
}
x.deltaRemove[key] = true
delete(x.deltaInsert, key)
}
}
}
// ToMap returns the map representation of the current set of variables.
//
// Mutating the returned map does not affect x.
func (x *Vars) ToMap() map[string]string {
snapshot := MergeMaps(x.base, x.deltaInsert)
for key := range x.deltaRemove {
delete(snapshot, key)
}
return snapshot
}
// ToSlice returns the slice representation of the current set of variables.
//
// Mutating the returned slice does not affect x.
func (x *Vars) ToSlice() []string {
return MapToSlice(x.ToMap())
}
// Base returns a copy of the original base environment.
//
// Mutating the returned map does not affect x.
func (x *Vars) Base() map[string]string {
return MergeMaps(x.base)
}
// Deltas returns the set of variables that have been mutated after
// initialization.
//
// If the last mutation for key K was Set or SetTokens, map[K] contains a
// non-nil pointer to the last value that was set. If the last mutation for key
// K was Delete, map[K] contains a nil pointer.
//
// Mutating the returned map does not affect x.
func (x *Vars) Deltas() map[string]*string {
deltas := make(map[string]*string, len(x.deltaInsert)+len(x.deltaRemove))
for key, value := range x.deltaInsert {
cp := value
deltas[key] = &cp
}
for key := range x.deltaRemove {
deltas[key] = nil
}
return deltas
}
// UpdateOS updates the OS with the current set of variables. All variables are
// visited in sorted order, and os.Setenv is called for each variable.
//
// Returns the first error encountered, if any.
func (x *Vars) UpdateOS() error {
var firstErr error
for _, kv := range x.ToSlice() {
key, value := SplitKeyValue(kv)
if err := os.Setenv(key, value); err != nil && firstErr == nil {
firstErr = err
}
}
return firstErr
}