yaml_map.go

Documentation: github.com/cli/go-gh/v2/internal/yamlmap

     1  // Package yamlmap is a wrapper of gopkg.in/yaml.v3 for interacting
     2  // with yaml data as if it were a map.
     3  package yamlmap
     4  
     5  import (
     6  	"errors"
     7  
     8  	"gopkg.in/yaml.v3"
     9  )
    10  
    11  const (
    12  	modified = "modifed"
    13  )
    14  
    15  type Map struct {
    16  	*yaml.Node
    17  }
    18  
    19  var ErrNotFound = errors.New("not found")
    20  var ErrInvalidYaml = errors.New("invalid yaml")
    21  var ErrInvalidFormat = errors.New("invalid format")
    22  
    23  func StringValue(value string) *Map {
    24  	return &Map{&yaml.Node{
    25  		Kind:  yaml.ScalarNode,
    26  		Tag:   "!!str",
    27  		Value: value,
    28  	}}
    29  }
    30  
    31  func MapValue() *Map {
    32  	return &Map{&yaml.Node{
    33  		Kind: yaml.MappingNode,
    34  		Tag:  "!!map",
    35  	}}
    36  }
    37  
    38  func NullValue() *Map {
    39  	return &Map{&yaml.Node{
    40  		Kind: yaml.ScalarNode,
    41  		Tag:  "!!null",
    42  	}}
    43  }
    44  
    45  func Unmarshal(data []byte) (*Map, error) {
    46  	var root yaml.Node
    47  	err := yaml.Unmarshal(data, &root)
    48  	if err != nil {
    49  		return nil, ErrInvalidYaml
    50  	}
    51  	if len(root.Content) == 0 {
    52  		return MapValue(), nil
    53  	}
    54  	if root.Content[0].Kind != yaml.MappingNode {
    55  		return nil, ErrInvalidFormat
    56  	}
    57  	return &Map{root.Content[0]}, nil
    58  }
    59  
    60  func Marshal(m *Map) ([]byte, error) {
    61  	return yaml.Marshal(m.Node)
    62  }
    63  
    64  func (m *Map) AddEntry(key string, value *Map) {
    65  	keyNode := &yaml.Node{
    66  		Kind:  yaml.ScalarNode,
    67  		Tag:   "!!str",
    68  		Value: key,
    69  	}
    70  	m.Content = append(m.Content, keyNode, value.Node)
    71  	m.SetModified()
    72  }
    73  
    74  func (m *Map) Empty() bool {
    75  	return m.Content == nil || len(m.Content) == 0
    76  }
    77  
    78  func (m *Map) FindEntry(key string) (*Map, error) {
    79  	// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
    80  	// When iterating over the content slice we only want to compare the keys of the yamlMap.
    81  	for i, v := range m.Content {
    82  		if i%2 != 0 {
    83  			continue
    84  		}
    85  		if v.Value == key {
    86  			if i+1 < len(m.Content) {
    87  				return &Map{m.Content[i+1]}, nil
    88  			}
    89  		}
    90  	}
    91  	return nil, ErrNotFound
    92  }
    93  
    94  func (m *Map) Keys() []string {
    95  	// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
    96  	// When iterating over the content slice we only want to select the keys of the yamlMap.
    97  	keys := []string{}
    98  	for i, v := range m.Content {
    99  		if i%2 != 0 {
   100  			continue
   101  		}
   102  		keys = append(keys, v.Value)
   103  	}
   104  	return keys
   105  }
   106  
   107  func (m *Map) RemoveEntry(key string) error {
   108  	// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
   109  	// When iterating over the content slice we only want to compare the keys of the yamlMap.
   110  	// If we find they key to remove, remove the key and its value from the content slice.
   111  	found, skipNext := false, false
   112  	newContent := []*yaml.Node{}
   113  	for i, v := range m.Content {
   114  		if skipNext {
   115  			skipNext = false
   116  			continue
   117  		}
   118  		if i%2 != 0 || v.Value != key {
   119  			newContent = append(newContent, v)
   120  		} else {
   121  			found = true
   122  			skipNext = true
   123  			m.SetModified()
   124  		}
   125  	}
   126  	if !found {
   127  		return ErrNotFound
   128  	}
   129  	m.Content = newContent
   130  	return nil
   131  }
   132  
   133  func (m *Map) SetEntry(key string, value *Map) {
   134  	// Note: The content slice of a yamlMap looks like [key1, value1, key2, value2, ...].
   135  	// When iterating over the content slice we only want to compare the keys of the yamlMap.
   136  	// If we find they key to set, set the next item in the content slice to the new value.
   137  	m.SetModified()
   138  	for i, v := range m.Content {
   139  		if i%2 != 0 || v.Value != key {
   140  			continue
   141  		}
   142  		if v.Value == key {
   143  			if i+1 < len(m.Content) {
   144  				m.Content[i+1] = value.Node
   145  				return
   146  			}
   147  		}
   148  	}
   149  	m.AddEntry(key, value)
   150  }
   151  
   152  // Note: This is a hack to introduce the concept of modified/unmodified
   153  // on top of gopkg.in/yaml.v3. This works by setting the Value property
   154  // of a MappingNode to a specific value and then later checking if the
   155  // node's Value property is that specific value. When a MappingNode gets
   156  // output as a string the Value property is not used, thus changing it
   157  // has no impact for our purposes.
   158  func (m *Map) SetModified() {
   159  	// Can not mark a non-mapping node as modified
   160  	if m.Node.Kind != yaml.MappingNode && m.Node.Tag == "!!null" {
   161  		m.Node.Kind = yaml.MappingNode
   162  		m.Node.Tag = "!!map"
   163  	}
   164  	if m.Node.Kind == yaml.MappingNode {
   165  		m.Node.Value = modified
   166  	}
   167  }
   168  
   169  // Traverse map using BFS to set all nodes as unmodified.
   170  func (m *Map) SetUnmodified() {
   171  	i := 0
   172  	queue := []*yaml.Node{m.Node}
   173  	for {
   174  		if i > (len(queue) - 1) {
   175  			break
   176  		}
   177  		q := queue[i]
   178  		i = i + 1
   179  		if q.Kind != yaml.MappingNode {
   180  			continue
   181  		}
   182  		q.Value = ""
   183  		queue = append(queue, q.Content...)
   184  	}
   185  }
   186  
   187  // Traverse map using BFS to searach for any nodes that have been modified.
   188  func (m *Map) IsModified() bool {
   189  	i := 0
   190  	queue := []*yaml.Node{m.Node}
   191  	for {
   192  		if i > (len(queue) - 1) {
   193  			break
   194  		}
   195  		q := queue[i]
   196  		i = i + 1
   197  		if q.Kind != yaml.MappingNode {
   198  			continue
   199  		}
   200  		if q.Value == modified {
   201  			return true
   202  		}
   203  		queue = append(queue, q.Content...)
   204  	}
   205  	return false
   206  }
   207  
   208  func (m *Map) String() string {
   209  	data, err := Marshal(m)
   210  	if err != nil {
   211  		return ""
   212  	}
   213  	return string(data)
   214  }
   215
View as plain text