|
1 | 1 | package storage |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "bytes" |
4 | 5 | "fmt" |
5 | 6 | "gopkg.in/yaml.v2" |
| 7 | + "io" |
6 | 8 | ) |
7 | 9 |
|
8 | 10 | // YAMLParser is a Parser implementation to handle yaml files. |
9 | 11 | type YAMLParser struct { |
10 | 12 | } |
11 | 13 |
|
| 14 | +// YAMLSingleDocument represents a decoded single-document YAML file. |
| 15 | +type YAMLSingleDocument struct { |
| 16 | + originalRequired |
| 17 | + Value interface{} |
| 18 | +} |
| 19 | + |
| 20 | +// RealValue returns the real value that dasel should use when processing data. |
| 21 | +func (d *YAMLSingleDocument) RealValue() interface{} { |
| 22 | + return d.Value |
| 23 | +} |
| 24 | + |
| 25 | +// YAMLMultiDocument represents a decoded multi-document YAML file. |
| 26 | +type YAMLMultiDocument struct { |
| 27 | + originalRequired |
| 28 | + Values []interface{} |
| 29 | +} |
| 30 | + |
| 31 | +// RealValue returns the real value that dasel should use when processing data. |
| 32 | +func (d *YAMLMultiDocument) RealValue() interface{} { |
| 33 | + return d.Values |
| 34 | +} |
| 35 | + |
12 | 36 | // FromBytes returns some Data that is represented by the given bytes. |
13 | 37 | func (p *YAMLParser) FromBytes(byteData []byte) (interface{}, error) { |
14 | | - var data interface{} |
15 | | - if err := yaml.Unmarshal(byteData, &data); err != nil { |
16 | | - return data, fmt.Errorf("could not unmarshal config data: %w", err) |
| 38 | + res := make([]interface{}, 0) |
| 39 | + |
| 40 | + decoder := yaml.NewDecoder(bytes.NewBuffer(byteData)) |
| 41 | + |
| 42 | +docLoop: |
| 43 | + for { |
| 44 | + var docData interface{} |
| 45 | + if err := decoder.Decode(&docData); err != nil { |
| 46 | + if err == io.EOF { |
| 47 | + break docLoop |
| 48 | + } |
| 49 | + return nil, fmt.Errorf("could not unmarshal config data: %w", err) |
| 50 | + } |
| 51 | + res = append(res, docData) |
| 52 | + } |
| 53 | + switch len(res) { |
| 54 | + case 0: |
| 55 | + return nil, nil |
| 56 | + case 1: |
| 57 | + return &YAMLSingleDocument{Value: res[0]}, nil |
| 58 | + default: |
| 59 | + return &YAMLMultiDocument{Values: res}, nil |
17 | 60 | } |
18 | | - return data, nil |
19 | 61 | } |
20 | 62 |
|
21 | 63 | // ToBytes returns a slice of bytes that represents the given value. |
22 | 64 | func (p *YAMLParser) ToBytes(value interface{}) ([]byte, error) { |
23 | | - return yaml.Marshal(value) |
| 65 | + buffer := new(bytes.Buffer) |
| 66 | + encoder := yaml.NewEncoder(buffer) |
| 67 | + defer encoder.Close() |
| 68 | + |
| 69 | + switch v := value.(type) { |
| 70 | + case *YAMLSingleDocument: |
| 71 | + if err := encoder.Encode(v.Value); err != nil { |
| 72 | + return nil, fmt.Errorf("could not encode single document: %w", err) |
| 73 | + } |
| 74 | + case *YAMLMultiDocument: |
| 75 | + for index, d := range v.Values { |
| 76 | + if err := encoder.Encode(d); err != nil { |
| 77 | + return nil, fmt.Errorf("could not encode multi document [%d]: %w", index, err) |
| 78 | + } |
| 79 | + } |
| 80 | + default: |
| 81 | + if err := encoder.Encode(v); err != nil { |
| 82 | + return nil, fmt.Errorf("could not encode default document type: %w", err) |
| 83 | + } |
| 84 | + } |
| 85 | + return buffer.Bytes(), nil |
24 | 86 | } |
0 commit comments