github.com/shuguocloud/go-zero@v1.3.0/core/mapping/yamlunmarshaler.go (about)

     1  package mapping
     2  
     3  import (
     4  	"encoding/json"
     5  	"errors"
     6  	"io"
     7  	"io/ioutil"
     8  
     9  	"gopkg.in/yaml.v2"
    10  )
    11  
    12  // To make .json & .yaml consistent, we just use json as the tag key.
    13  const yamlTagKey = "json"
    14  
    15  var (
    16  	// ErrUnsupportedType is an error that indicates the config format is not supported.
    17  	ErrUnsupportedType = errors.New("only map-like configs are suported")
    18  
    19  	yamlUnmarshaler = NewUnmarshaler(yamlTagKey)
    20  )
    21  
    22  // UnmarshalYamlBytes unmarshals content into v.
    23  func UnmarshalYamlBytes(content []byte, v interface{}) error {
    24  	return unmarshalYamlBytes(content, v, yamlUnmarshaler)
    25  }
    26  
    27  // UnmarshalYamlReader unmarshals content from reader into v.
    28  func UnmarshalYamlReader(reader io.Reader, v interface{}) error {
    29  	return unmarshalYamlReader(reader, v, yamlUnmarshaler)
    30  }
    31  
    32  func unmarshalYamlBytes(content []byte, v interface{}, unmarshaler *Unmarshaler) error {
    33  	var o interface{}
    34  	if err := yamlUnmarshal(content, &o); err != nil {
    35  		return err
    36  	}
    37  
    38  	if m, ok := o.(map[string]interface{}); ok {
    39  		return unmarshaler.Unmarshal(m, v)
    40  	}
    41  
    42  	return ErrUnsupportedType
    43  }
    44  
    45  func unmarshalYamlReader(reader io.Reader, v interface{}, unmarshaler *Unmarshaler) error {
    46  	content, err := ioutil.ReadAll(reader)
    47  	if err != nil {
    48  		return err
    49  	}
    50  
    51  	return unmarshalYamlBytes(content, v, unmarshaler)
    52  }
    53  
    54  // yamlUnmarshal YAML to map[string]interface{} instead of map[interface{}]interface{}.
    55  func yamlUnmarshal(in []byte, out interface{}) error {
    56  	var res interface{}
    57  	if err := yaml.Unmarshal(in, &res); err != nil {
    58  		return err
    59  	}
    60  
    61  	*out.(*interface{}) = cleanupMapValue(res)
    62  	return nil
    63  }
    64  
    65  func cleanupInterfaceMap(in map[interface{}]interface{}) map[string]interface{} {
    66  	res := make(map[string]interface{})
    67  	for k, v := range in {
    68  		res[Repr(k)] = cleanupMapValue(v)
    69  	}
    70  	return res
    71  }
    72  
    73  func cleanupInterfaceNumber(in interface{}) json.Number {
    74  	return json.Number(Repr(in))
    75  }
    76  
    77  func cleanupInterfaceSlice(in []interface{}) []interface{} {
    78  	res := make([]interface{}, len(in))
    79  	for i, v := range in {
    80  		res[i] = cleanupMapValue(v)
    81  	}
    82  	return res
    83  }
    84  
    85  func cleanupMapValue(v interface{}) interface{} {
    86  	switch v := v.(type) {
    87  	case []interface{}:
    88  		return cleanupInterfaceSlice(v)
    89  	case map[interface{}]interface{}:
    90  		return cleanupInterfaceMap(v)
    91  	case bool, string:
    92  		return v
    93  	case int, uint, int8, uint8, int16, uint16, int32, uint32, int64, uint64, float32, float64:
    94  		return cleanupInterfaceNumber(v)
    95  	default:
    96  		return Repr(v)
    97  	}
    98  }