From 3a285e02264bd366a585bfd7d735c58d869a6c12 Mon Sep 17 00:00:00 2001
From: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
Date: Fri, 21 Jun 2024 14:27:56 +0200
Subject: [PATCH] feat(encoding): expose new encoding layer

Signed-off-by: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
---
 encoding.go   | 113 ++++++++++++++++++++++++++++++++++++++++++++++++++
 viper.go      |  23 +++++++++-
 viper_test.go |   6 +--
 3 files changed, 137 insertions(+), 5 deletions(-)
 create mode 100644 encoding.go

diff --git a/encoding.go b/encoding.go
new file mode 100644
index 0000000..7d48d0d
--- /dev/null
+++ b/encoding.go
@@ -0,0 +1,113 @@
+package viper
+
+import (
+	"github.com/spf13/viper/internal/encoding/dotenv"
+	"github.com/spf13/viper/internal/encoding/hcl"
+	"github.com/spf13/viper/internal/encoding/ini"
+	"github.com/spf13/viper/internal/encoding/javaproperties"
+	"github.com/spf13/viper/internal/encoding/json"
+	"github.com/spf13/viper/internal/encoding/toml"
+	"github.com/spf13/viper/internal/encoding/yaml"
+)
+
+// Encoder encodes Viper's internal data structures into a byte representation.
+// It's primarily used for encoding a map[string]any into a file format.
+type Encoder interface {
+	Encode(v map[string]any) ([]byte, error)
+}
+
+// Decoder decodes the contents of a byte slice into Viper's internal data structures.
+// It's primarily used for decoding contents of a file into a map[string]any.
+type Decoder interface {
+	Decode(b []byte, v map[string]any) error
+}
+
+// Codec combines [Encoder] and [Decoder] interfaces.
+type Codec interface {
+	Encoder
+	Decoder
+}
+
+// EncoderRegistry returns an [Encoder] for a given format.
+// The second return value is false if no [Encoder] is registered for the format.
+type EncoderRegistry interface {
+	Encoder(format string) (Encoder, bool)
+}
+
+// DecoderRegistry returns an [Decoder] for a given format.
+// The second return value is false if no [Decoder] is registered for the format.
+type DecoderRegistry interface {
+	Decoder(format string) (Decoder, bool)
+}
+
+// [CodecRegistry] combines [EncoderRegistry] and [DecoderRegistry] interfaces.
+type CodecRegistry interface {
+	EncoderRegistry
+	DecoderRegistry
+}
+
+// WithEncoderRegistry sets a custom [EncoderRegistry].
+func WithEncoderRegistry(r EncoderRegistry) Option {
+	return optionFunc(func(v *Viper) {
+		v.encoderRegistry2 = r
+	})
+}
+
+// WithDecoderRegistry sets a custom [DecoderRegistry].
+func WithDecoderRegistry(r DecoderRegistry) Option {
+	return optionFunc(func(v *Viper) {
+		v.decoderRegistry2 = r
+	})
+}
+
+// WithCodecRegistry sets a custom [EncoderRegistry] and [DecoderRegistry].
+func WithCodecRegistry(r CodecRegistry) Option {
+	return optionFunc(func(v *Viper) {
+		v.encoderRegistry2 = r
+		v.decoderRegistry2 = r
+	})
+}
+
+type codecRegistry struct {
+	v *Viper
+}
+
+func (r codecRegistry) Encoder(format string) (Encoder, bool) {
+	return r.codec(format)
+}
+
+func (r codecRegistry) Decoder(format string) (Decoder, bool) {
+	return r.codec(format)
+}
+
+func (r codecRegistry) codec(format string) (Codec, bool) {
+	switch format {
+	case "yaml", "yml":
+		return yaml.Codec{}, true
+
+	case "json":
+		return json.Codec{}, true
+
+	case "toml":
+		return toml.Codec{}, true
+
+	case "hcl", "tfvars":
+		return hcl.Codec{}, true
+
+	case "ini":
+		return ini.Codec{
+			KeyDelimiter: r.v.keyDelim,
+			LoadOptions:  r.v.iniLoadOptions,
+		}, true
+
+	case "properties", "props", "prop": // Note: This breaks writing a properties file.
+		return &javaproperties.Codec{
+			KeyDelimiter: v.keyDelim,
+		}, true
+
+	case "dotenv", "env":
+		return &dotenv.Codec{}, true
+	}
+
+	return nil, false
+}
diff --git a/viper.go b/viper.go
index 30ae4aa..151ec2d 100644
--- a/viper.go
+++ b/viper.go
@@ -195,6 +195,9 @@ type Viper struct {
 	encoderRegistry *encoding.EncoderRegistry
 	decoderRegistry *encoding.DecoderRegistry
 
+	encoderRegistry2 EncoderRegistry
+	decoderRegistry2 DecoderRegistry
+
 	experimentalFinder     bool
 	experimentalBindStruct bool
 }
@@ -217,6 +220,11 @@ func New() *Viper {
 	v.typeByDefValue = false
 	v.logger = slog.New(&discardHandler{})
 
+	codecRegistry := codecRegistry{v: v}
+
+	v.encoderRegistry2 = codecRegistry
+	v.decoderRegistry2 = codecRegistry
+
 	v.resetEncoding()
 
 	v.experimentalFinder = features.Finder
@@ -1715,7 +1723,12 @@ func (v *Viper) unmarshalReader(in io.Reader, c map[string]any) error {
 
 	switch format := strings.ToLower(v.getConfigType()); format {
 	case "yaml", "yml", "json", "toml", "hcl", "tfvars", "ini", "properties", "props", "prop", "dotenv", "env":
-		err := v.decoderRegistry.Decode(format, buf.Bytes(), c)
+		decoder, ok := v.decoderRegistry2.Decoder(format)
+		if !ok {
+			return ConfigParseError{errors.New("decoder not found")}
+		}
+
+		err := decoder.Decode(buf.Bytes(), c)
 		if err != nil {
 			return ConfigParseError{err}
 		}
@@ -1730,7 +1743,13 @@ func (v *Viper) marshalWriter(f afero.File, configType string) error {
 	c := v.AllSettings()
 	switch configType {
 	case "yaml", "yml", "json", "toml", "hcl", "tfvars", "ini", "prop", "props", "properties", "dotenv", "env":
-		b, err := v.encoderRegistry.Encode(configType, c)
+		encoder, ok := v.encoderRegistry2.Encoder(configType)
+		if !ok {
+			// TODO: return a proper error
+			return ConfigMarshalError{errors.New("encoder not found")}
+		}
+
+		b, err := encoder.Encode(c)
 		if err != nil {
 			return ConfigMarshalError{err}
 		}
diff --git a/viper_test.go b/viper_test.go
index 2edc83f..05ca0ce 100644
--- a/viper_test.go
+++ b/viper_test.go
@@ -1849,11 +1849,11 @@ var jsonWriteExpected = []byte(`{
   "type": "donut"
 }`)
 
-var propertiesWriteExpected = []byte(`p_id = 0001
-p_type = donut
+var propertiesWriteExpected = []byte(`p_batters.batter.type = Regular
+p_id = 0001
 p_name = Cake
 p_ppu = 0.55
-p_batters.batter.type = Regular
+p_type = donut
 `)
 
 // var yamlWriteExpected = []byte(`age: 35