From 15ad72deb21199dde5877cc5ada6e167607a93e7 Mon Sep 17 00:00:00 2001
From: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
Date: Mon, 24 Jun 2024 11:37:05 +0200
Subject: [PATCH] feat(encoding): add default codec registry

Signed-off-by: Mark Sagi-Kazar <mark.sagikazar@gmail.com>
---
 encoding.go      | 84 ++++++++++++++++++++++++++++++++++++++++++++++++
 encoding_test.go | 39 ++++++++++++++++++++++
 2 files changed, 123 insertions(+)
 create mode 100644 encoding_test.go

diff --git a/encoding.go b/encoding.go
index 2563348..d4f38f6 100644
--- a/encoding.go
+++ b/encoding.go
@@ -1,6 +1,9 @@
 package viper
 
 import (
+	"strings"
+	"sync"
+
 	"github.com/spf13/viper/internal/encoding/dotenv"
 	"github.com/spf13/viper/internal/encoding/hcl"
 	"github.com/spf13/viper/internal/encoding/ini"
@@ -137,3 +140,84 @@ func (r codecRegistry) codec(format string) (Codec, bool) {
 
 	return nil, false
 }
+
+// DefaultCodecRegistry
+type DefaultCodecRegistry struct {
+	codecs map[string]Codec
+
+	mu   sync.RWMutex
+	once sync.Once
+}
+
+// NewCodecRegistry returns a new [CodecRegistry], ready to accept custom [Codec]s.
+func NewCodecRegistry() *DefaultCodecRegistry {
+	r := &DefaultCodecRegistry{}
+
+	r.init()
+
+	return r
+}
+
+func (r *DefaultCodecRegistry) init() {
+	r.once.Do(func() {
+		r.codecs = map[string]Codec{}
+	})
+}
+
+// RegisterCodec registers a custom [Codec].
+func (r *DefaultCodecRegistry) RegisterCodec(format string, codec Codec) error {
+	r.init()
+
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	r.codecs[strings.ToLower(format)] = codec
+
+	return nil
+}
+
+func (r *DefaultCodecRegistry) Encoder(format string) (Encoder, error) {
+	encoder, ok := r.codec(format)
+	if !ok {
+		return nil, ErrEncoderNotFound
+	}
+
+	return encoder, nil
+}
+
+func (r *DefaultCodecRegistry) Decoder(format string) (Decoder, error) {
+	decoder, ok := r.codec(format)
+	if !ok {
+		return nil, ErrDecoderNotFound
+	}
+
+	return decoder, nil
+}
+
+func (r *DefaultCodecRegistry) codec(format string) (Codec, bool) {
+	r.mu.Lock()
+	defer r.mu.Unlock()
+
+	if r.codecs != nil {
+		codec, ok := r.codecs[format]
+		if ok {
+			return codec, true
+		}
+	}
+
+	switch format {
+	case "yaml", "yml":
+		return yaml.Codec{}, true
+
+	case "json":
+		return json.Codec{}, true
+
+	case "toml":
+		return toml.Codec{}, true
+
+	case "dotenv", "env":
+		return &dotenv.Codec{}, true
+	}
+
+	return nil, false
+}
diff --git a/encoding_test.go b/encoding_test.go
new file mode 100644
index 0000000..61199c0
--- /dev/null
+++ b/encoding_test.go
@@ -0,0 +1,39 @@
+package viper
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+type codec struct{}
+
+func (codec) Encode(_ map[string]any) ([]byte, error) {
+	return nil, nil
+}
+
+func (codec) Decode(_ []byte, _ map[string]any) error {
+	return nil
+}
+
+func TestDefaultCodecRegistry(t *testing.T) {
+	t.Run("OK", func(t *testing.T) {
+		registry := NewCodecRegistry()
+
+		c := codec{}
+
+		err := registry.RegisterCodec("myformat", c)
+		require.NoError(t, err)
+
+		encoder, err := registry.Encoder("myformat")
+		require.NoError(t, err)
+
+		assert.Equal(t, c, encoder)
+
+		decoder, err := registry.Decoder("myformat")
+		require.NoError(t, err)
+
+		assert.Equal(t, c, decoder)
+	})
+}