1
package viper
2

3
import (
4
	"errors"
5
	"strings"
6
	"sync"
7

8
	"github.com/spf13/viper/internal/encoding/dotenv"
9
	"github.com/spf13/viper/internal/encoding/json"
10
	"github.com/spf13/viper/internal/encoding/toml"
11
	"github.com/spf13/viper/internal/encoding/yaml"
12
)
13

14
// Encoder encodes Viper's internal data structures into a byte representation.
15
// It's primarily used for encoding a map[string]any into a file format.
16
type Encoder interface {
17
	Encode(v map[string]any) ([]byte, error)
18
}
19

20
// Decoder decodes the contents of a byte slice into Viper's internal data structures.
21
// It's primarily used for decoding contents of a file into a map[string]any.
22
type Decoder interface {
23
	Decode(b []byte, v map[string]any) error
24
}
25

26
// Codec combines [Encoder] and [Decoder] interfaces.
27
type Codec interface {
28
	Encoder
29
	Decoder
30
}
31

32
// TODO: consider adding specific errors for not found scenarios
33

34
// EncoderRegistry returns an [Encoder] for a given format.
35
//
36
// Format is case-insensitive.
37
//
38
// [EncoderRegistry] returns an error if no [Encoder] is registered for the format.
39
type EncoderRegistry interface {
40
	Encoder(format string) (Encoder, error)
41
}
42

43
// DecoderRegistry returns an [Decoder] for a given format.
44
//
45
// Format is case-insensitive.
46
//
47
// [DecoderRegistry] returns an error if no [Decoder] is registered for the format.
48
type DecoderRegistry interface {
49
	Decoder(format string) (Decoder, error)
50
}
51

52
// [CodecRegistry] combines [EncoderRegistry] and [DecoderRegistry] interfaces.
53
type CodecRegistry interface {
54
	EncoderRegistry
55
	DecoderRegistry
56
}
57

58
// WithEncoderRegistry sets a custom [EncoderRegistry].
59
func WithEncoderRegistry(r EncoderRegistry) Option {
60
	return optionFunc(func(v *Viper) {
61
		if r == nil {
62
			return
63
		}
64

65
		v.encoderRegistry = r
66
	})
67
}
68

69
// WithDecoderRegistry sets a custom [DecoderRegistry].
70
func WithDecoderRegistry(r DecoderRegistry) Option {
71
	return optionFunc(func(v *Viper) {
72
		if r == nil {
73
			return
74
		}
75

76
		v.decoderRegistry = r
77
	})
78
}
79

80
// WithCodecRegistry sets a custom [EncoderRegistry] and [DecoderRegistry].
81
func WithCodecRegistry(r CodecRegistry) Option {
82
	return optionFunc(func(v *Viper) {
83
		if r == nil {
84
			return
85
		}
86

87
		v.encoderRegistry = r
88
		v.decoderRegistry = r
89
	})
90
}
91

92
// DefaultCodecRegistry is a simple implementation of [CodecRegistry] that allows registering custom [Codec]s.
93
type DefaultCodecRegistry struct {
94
	codecs map[string]Codec
95

96
	mu   sync.RWMutex
97
	once sync.Once
98
}
99

100
// NewCodecRegistry returns a new [CodecRegistry], ready to accept custom [Codec]s.
101
func NewCodecRegistry() *DefaultCodecRegistry {
102
	r := &DefaultCodecRegistry{}
103

104
	r.init()
105

106
	return r
107
}
108

109
func (r *DefaultCodecRegistry) init() {
110
	r.once.Do(func() {
111
		r.codecs = map[string]Codec{}
112
	})
113
}
114

115
// RegisterCodec registers a custom [Codec].
116
//
117
// Format is case-insensitive.
118
func (r *DefaultCodecRegistry) RegisterCodec(format string, codec Codec) error {
119
	r.init()
120

121
	r.mu.Lock()
122
	defer r.mu.Unlock()
123

124
	r.codecs[strings.ToLower(format)] = codec
125

126
	return nil
127
}
128

129
// Encoder implements the [EncoderRegistry] interface.
130
//
131
// Format is case-insensitive.
132
func (r *DefaultCodecRegistry) Encoder(format string) (Encoder, error) {
133
	encoder, ok := r.codec(format)
134
	if !ok {
135
		return nil, errors.New("encoder not found for this format")
136
	}
137

138
	return encoder, nil
139
}
140

141
// Decoder implements the [DecoderRegistry] interface.
142
//
143
// Format is case-insensitive.
144
func (r *DefaultCodecRegistry) Decoder(format string) (Decoder, error) {
145
	decoder, ok := r.codec(format)
146
	if !ok {
147
		return nil, errors.New("decoder not found for this format")
148
	}
149

150
	return decoder, nil
151
}
152

153
func (r *DefaultCodecRegistry) codec(format string) (Codec, bool) {
154
	r.mu.Lock()
155
	defer r.mu.Unlock()
156

157
	format = strings.ToLower(format)
158

159
	if r.codecs != nil {
160
		codec, ok := r.codecs[format]
161
		if ok {
162
			return codec, true
163
		}
164
	}
165

166
	switch format {
167
	case "yaml", "yml":
168
		return yaml.Codec{}, true
169

170
	case "json":
171
		return json.Codec{}, true
172

173
	case "toml":
174
		return toml.Codec{}, true
175

176
	case "dotenv", "env":
177
		return &dotenv.Codec{}, true
178
	}
179

180
	return nil, false
181
}
182

183