github.com/tomwright/dasel@v1.27.3/storage/yaml.go (about)

     1  package storage
     2  
     3  import (
     4  	"bytes"
     5  	"fmt"
     6  	"gopkg.in/yaml.v2"
     7  	"io"
     8  )
     9  
    10  func init() {
    11  	registerReadParser([]string{"yaml", "yml"}, []string{".yaml", ".yml"}, &YAMLParser{})
    12  	registerWriteParser([]string{"yaml", "yml"}, []string{".yaml", ".yml"}, &YAMLParser{})
    13  }
    14  
    15  // YAMLParser is a Parser implementation to handle yaml files.
    16  type YAMLParser struct {
    17  }
    18  
    19  // FromBytes returns some data that is represented by the given bytes.
    20  func (p *YAMLParser) FromBytes(byteData []byte) (interface{}, error) {
    21  	res := make([]interface{}, 0)
    22  
    23  	decoder := yaml.NewDecoder(bytes.NewBuffer(byteData))
    24  
    25  docLoop:
    26  	for {
    27  		var docData interface{}
    28  		if err := decoder.Decode(&docData); err != nil {
    29  			if err == io.EOF {
    30  				break docLoop
    31  			}
    32  			return nil, fmt.Errorf("could not unmarshal data: %w", err)
    33  		}
    34  
    35  		formattedDocData := cleanupYamlMapValue(docData)
    36  
    37  		res = append(res, formattedDocData)
    38  	}
    39  	switch len(res) {
    40  	case 0:
    41  		return nil, nil
    42  	case 1:
    43  		return &BasicSingleDocument{Value: res[0]}, nil
    44  	default:
    45  		return &BasicMultiDocument{Values: res}, nil
    46  	}
    47  }
    48  
    49  func cleanupYamlInterfaceArray(in []interface{}) []interface{} {
    50  	res := make([]interface{}, len(in))
    51  	for i, v := range in {
    52  		res[i] = cleanupYamlMapValue(v)
    53  	}
    54  	return res
    55  }
    56  
    57  func cleanupYamlInterfaceMap(in map[interface{}]interface{}) map[string]interface{} {
    58  	res := make(map[string]interface{})
    59  	for k, v := range in {
    60  		res[fmt.Sprint(k)] = cleanupYamlMapValue(v)
    61  	}
    62  	return res
    63  }
    64  
    65  func cleanupYamlMapValue(v interface{}) interface{} {
    66  	switch v := v.(type) {
    67  	case []interface{}:
    68  		return cleanupYamlInterfaceArray(v)
    69  	case map[interface{}]interface{}:
    70  		return cleanupYamlInterfaceMap(v)
    71  	case string:
    72  		return v
    73  	default:
    74  		return v
    75  	}
    76  }
    77  
    78  // ToBytes returns a slice of bytes that represents the given value.
    79  func (p *YAMLParser) ToBytes(value interface{}, options ...ReadWriteOption) ([]byte, error) {
    80  	buffer := new(bytes.Buffer)
    81  	encoder := yaml.NewEncoder(buffer)
    82  	defer encoder.Close()
    83  
    84  	colourise := false
    85  
    86  	for _, o := range options {
    87  		switch o.Key {
    88  		case OptionColourise:
    89  			if value, ok := o.Value.(bool); ok {
    90  				colourise = value
    91  			}
    92  		}
    93  	}
    94  
    95  	switch v := value.(type) {
    96  	case SingleDocument:
    97  		if err := encoder.Encode(v.Document()); err != nil {
    98  			return nil, fmt.Errorf("could not encode single document: %w", err)
    99  		}
   100  	case MultiDocument:
   101  		for index, d := range v.Documents() {
   102  			if err := encoder.Encode(d); err != nil {
   103  				return nil, fmt.Errorf("could not encode multi document [%d]: %w", index, err)
   104  			}
   105  		}
   106  	default:
   107  		if err := encoder.Encode(v); err != nil {
   108  			return nil, fmt.Errorf("could not encode default document type: %w", err)
   109  		}
   110  	}
   111  
   112  	if colourise {
   113  		if err := ColouriseBuffer(buffer, "yaml"); err != nil {
   114  			return nil, fmt.Errorf("could not colourise output: %w", err)
   115  		}
   116  	}
   117  
   118  	return buffer.Bytes(), nil
   119  }