confkit/flags_source.go at main · MimoJanra/confkit · GitHub

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
package confkit

import (
	"context"
	"os"
	"strings"
	"sync"

	"github.com/MimoJanra/confkit/structtags"
)

func FromFlags() Source {
	return &flagsSource{args: os.Args[1:]}
}

func FromFlagsWithArgs(args []string) Source {
	return &flagsSource{args: args}
}

type flagsSource struct {
	args   []string
	parsed map[string]string
	once   sync.Once
}

func (f *flagsSource) Name() string {
	return "flag"
}

func (f *flagsSource) parseArgs() {
	f.once.Do(func() {
		f.parsed = make(map[string]string)
		args := f.args
		for i := 0; i < len(args); i++ {
			arg := args[i]
			if strings.HasPrefix(arg, "--") {
				name := arg[2:]
				if eq := strings.Index(name, "="); eq != -1 {
					appendFlag(f.parsed, name[:eq], name[eq+1:])
				} else if i+1 < len(args) && !isFlagToken(args[i+1]) {
					appendFlag(f.parsed, name, args[i+1])
					i++
				} else {
					appendFlag(f.parsed, name, "true")
				}
			} else if strings.HasPrefix(arg, "-") && len(arg) == 2 {
				name := arg[1:]
				if i+1 < len(args) && !isFlagToken(args[i+1]) {
					appendFlag(f.parsed, name, args[i+1])
					i++
				} else {
					appendFlag(f.parsed, name, "true")
				}
			} else if strings.HasPrefix(arg, "-") && len(arg) > 2 && arg[2] == '=' {
				name := arg[1:2]
				appendFlag(f.parsed, name, arg[3:])
			}
		}
	})
}

func isFlagToken(s string) bool {
	if !strings.HasPrefix(s, "-") {
		return false
	}
	if len(s) > 1 && (s[1] >= '0' && s[1] <= '9' || s[1] == '.') {
		return false
	}
	return true
}

func appendFlag(m map[string]string, key, val string) {
	if existing, ok := m[key]; ok {
		m[key] = existing + "," + val
	} else {
		m[key] = val
	}
}

func (f *flagsSource) Lookup(_ context.Context, field *FieldInfo) (any, bool, error) {
	f.parseArgs()

	flagName := field.Tags["flag"]
	if flagName == "" {
		flagName = toKebabCase(field.Name)
	}

	if val, ok := f.parsed[flagName]; ok {
		return val, true, nil
	}

	snakeName := structtags.SnakeCase(field.Name)
	if snakeName != flagName {
		if val, ok := f.parsed[snakeName]; ok {
			return val, true, nil
		}
	}

	if short := field.Tags["short"]; short != "" {
		if val, ok := f.parsed[short]; ok {
			return val, true, nil
		}
	}

	return "", false, nil
}

func toKebabCase(s string) string {
	snake := structtags.SnakeCase(s)
	return strings.ReplaceAll(snake, "_", "-")
}