From e989bb04e387a921fc45866a35bd3f2462db87b0 Mon Sep 17 00:00:00 2001
From: Tony Chen <zhihui_chen@foxmail.com>
Date: Mon, 17 May 2021 21:52:23 +0800
Subject: [PATCH] add yaml encoder (#924)

---
 config/config.go        | 14 ++++++++++++--
 config/file/file.go     |  5 +++--
 config/file/format.go   | 10 ++++++++++
 config/source.go        |  6 +++---
 encoding/yaml/yaml.go   | 37 +++++++++++++++++++++++++++++++++++++
 examples/config/main.go |  4 ----
 go.mod                  |  1 +
 go.sum                  |  2 ++
 transport/transport.go  |  1 +
 9 files changed, 69 insertions(+), 11 deletions(-)
 create mode 100644 config/file/format.go
 create mode 100644 encoding/yaml/yaml.go

diff --git a/config/config.go b/config/config.go
index de08bbdb7..f078a1008 100644
--- a/config/config.go
+++ b/config/config.go
@@ -1,13 +1,20 @@
 package config
 
 import (
-	"encoding/json"
 	"errors"
+	"fmt"
 	"reflect"
 	"sync"
 	"time"
 
+	"github.com/go-kratos/kratos/v2/encoding"
 	"github.com/go-kratos/kratos/v2/log"
+
+	// init encoding
+	_ "github.com/go-kratos/kratos/v2/encoding/json"
+	_ "github.com/go-kratos/kratos/v2/encoding/proto"
+	_ "github.com/go-kratos/kratos/v2/encoding/xml"
+	_ "github.com/go-kratos/kratos/v2/encoding/yaml"
 )
 
 var (
@@ -45,7 +52,10 @@ func New(opts ...Option) Config {
 	options := options{
 		logger: log.DefaultLogger,
 		decoder: func(kv *KeyValue, v map[string]interface{}) error {
-			return json.Unmarshal(kv.Value, &v)
+			if codec := encoding.GetCodec(kv.Format); codec != nil {
+				return codec.Unmarshal(kv.Value, &v)
+			}
+			return fmt.Errorf("unsupported key: %s format: %s", kv.Key, kv.Format)
 		},
 	}
 	for _, o := range opts {
diff --git a/config/file/file.go b/config/file/file.go
index 5959666f0..9a918c088 100644
--- a/config/file/file.go
+++ b/config/file/file.go
@@ -35,8 +35,9 @@ func (f *file) loadFile(path string) (*config.KeyValue, error) {
 		return nil, err
 	}
 	return &config.KeyValue{
-		Key:   info.Name(),
-		Value: data,
+		Key:    info.Name(),
+		Format: format(info.Name()),
+		Value:  data,
 	}, nil
 }
 
diff --git a/config/file/format.go b/config/file/format.go
new file mode 100644
index 000000000..9e20e60b1
--- /dev/null
+++ b/config/file/format.go
@@ -0,0 +1,10 @@
+package file
+
+import "strings"
+
+func format(name string) string {
+	if p := strings.Split(name, "."); len(p) > 1 {
+		return p[len(p)-1]
+	}
+	return ""
+}
diff --git a/config/source.go b/config/source.go
index 225a162ae..bd86e4d2b 100644
--- a/config/source.go
+++ b/config/source.go
@@ -2,9 +2,9 @@ package config
 
 // KeyValue is config key value.
 type KeyValue struct {
-	Key      string
-	Value    []byte
-	Metadata map[string]string
+	Key    string
+	Value  []byte
+	Format string
 }
 
 // Source is config source.
diff --git a/encoding/yaml/yaml.go b/encoding/yaml/yaml.go
new file mode 100644
index 000000000..4e49284b2
--- /dev/null
+++ b/encoding/yaml/yaml.go
@@ -0,0 +1,37 @@
+package yaml
+
+import (
+	"reflect"
+
+	"github.com/go-kratos/kratos/v2/encoding"
+	"gopkg.in/yaml.v2"
+)
+
+// Name is the name registered for the json codec.
+const Name = "yaml"
+
+func init() {
+	encoding.RegisterCodec(codec{})
+}
+
+// codec is a Codec implementation with json.
+type codec struct{}
+
+func (codec) Marshal(v interface{}) ([]byte, error) {
+	return yaml.Marshal(v)
+}
+
+func (codec) Unmarshal(data []byte, v interface{}) error {
+	rv := reflect.ValueOf(v)
+	for rv.Kind() == reflect.Ptr {
+		if rv.IsNil() {
+			rv.Set(reflect.New(rv.Type().Elem()))
+		}
+		rv = rv.Elem()
+	}
+	return yaml.Unmarshal(data, v)
+}
+
+func (codec) Name() string {
+	return Name
+}
diff --git a/examples/config/main.go b/examples/config/main.go
index f93bb993a..fab56b1c6 100644
--- a/examples/config/main.go
+++ b/examples/config/main.go
@@ -6,7 +6,6 @@ import (
 
 	"github.com/go-kratos/kratos/v2/config"
 	"github.com/go-kratos/kratos/v2/config/file"
-	"gopkg.in/yaml.v2"
 )
 
 var flagconf string
@@ -21,9 +20,6 @@ func main() {
 		config.WithSource(
 			file.NewSource(flagconf),
 		),
-		config.WithDecoder(func(kv *config.KeyValue, v map[string]interface{}) error {
-			return yaml.Unmarshal(kv.Value, v)
-		}),
 	)
 	if err := c.Load(); err != nil {
 		panic(err)
diff --git a/go.mod b/go.mod
index 8dd913c06..2a412fc40 100644
--- a/go.mod
+++ b/go.mod
@@ -18,4 +18,5 @@ require (
 	google.golang.org/genproto v0.0.0-20210114201628-6edceaf6022f
 	google.golang.org/grpc v1.35.0
 	google.golang.org/protobuf v1.25.0
+	gopkg.in/yaml.v2 v2.4.0 // indirect
 )
diff --git a/go.sum b/go.sum
index 06a9de0a5..674faf2e3 100644
--- a/go.sum
+++ b/go.sum
@@ -125,6 +125,8 @@ gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405 h1:yhCVgyC4o1eVCa2tZl7eS0r+
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/yaml.v2 v2.2.2 h1:ZCJp+EgiOT7lHqUV2J862kp8Qj64Jo6az82+3Td9dZw=
 gopkg.in/yaml.v2 v2.2.2/go.mod h1:hI93XBmqTisBFMUTm0b8Fm+jr3Dg1NNxqwp+5A1VGuI=
+gopkg.in/yaml.v2 v2.4.0 h1:D8xgwECY7CYvx+Y2n4sBz93Jn9JRvxdiyyo8CTfuKaY=
+gopkg.in/yaml.v2 v2.4.0/go.mod h1:RDklbk79AGWmwhnvt/jBztapEOGDOx6ZbXqjP6csGnQ=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c h1:dUUwHk2QECo/6vqA44rthZ8ie2QXMNeKRTHCNY2nXvo=
 gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM=
 honnef.co/go/tools v0.0.0-20190102054323-c2f93a96b099/go.mod h1:rf3lG4BRIbNafJWhAfAdb/ePZxsR/4RtNHQocxwk9r4=
diff --git a/transport/transport.go b/transport/transport.go
index b268aaefa..f3eee1ddb 100644
--- a/transport/transport.go
+++ b/transport/transport.go
@@ -7,6 +7,7 @@ import (
 	_ "github.com/go-kratos/kratos/v2/encoding/json"
 	_ "github.com/go-kratos/kratos/v2/encoding/proto"
 	_ "github.com/go-kratos/kratos/v2/encoding/xml"
+	_ "github.com/go-kratos/kratos/v2/encoding/yaml"
 )
 
 // Server is transport server.