talm

1
// Package yamltools provides functions for handling YAML nodes, such as copying comments, applying comments,
2
// and diffing YAML documents.
3
package yamltools
4

5
import (
6
	"bytes"
7
	"strings"
8

9
	"gopkg.in/yaml.v3"
10
)
11

12
// CopyComments updates the comments in dstNode considering the structure of whitespace.
13
func CopyComments(srcNode, dstNode *yaml.Node, path string, dstPaths map[string]*yaml.Node) {
14
	if srcNode.HeadComment != "" || srcNode.LineComment != "" || srcNode.FootComment != "" {
15
		dstPaths[path] = srcNode
16
	}
17

18
	for i := 0; i < len(srcNode.Content); i++ {
19
		newPath := path + "/" + srcNode.Content[i].Value
20
		if srcNode.Kind == yaml.SequenceNode {
21
			newPath = path + "/" + string(i)
22
		}
23
		CopyComments(srcNode.Content[i], dstNode, newPath, dstPaths)
24
	}
25
}
26

27
// ApplyComments applies the copied comments to the target document.
28
func ApplyComments(dstNode *yaml.Node, path string, dstPaths map[string]*yaml.Node) {
29
	if srcNode, ok := dstPaths[path]; ok {
30
		dstNode.HeadComment = mergeComments(dstNode.HeadComment, srcNode.HeadComment)
31
		dstNode.LineComment = mergeComments(dstNode.LineComment, srcNode.LineComment)
32
		dstNode.FootComment = mergeComments(dstNode.FootComment, srcNode.FootComment)
33
	}
34

35
	for i := 0; i < len(dstNode.Content); i++ {
36
		newPath := path + "/" + dstNode.Content[i].Value
37
		if dstNode.Kind == yaml.SequenceNode {
38
			newPath = path + "/" + string(i)
39
		}
40
		ApplyComments(dstNode.Content[i], newPath, dstPaths)
41
	}
42
}
43

44
// mergeComments combines old and new comments considering empty lines.
45
func mergeComments(oldComment, newComment string) string {
46
	if oldComment == "" {
47
		return newComment
48
	}
49
	if newComment == "" {
50
		return oldComment
51
	}
52
	return strings.TrimSpace(oldComment) + "\n\n" + strings.TrimSpace(newComment)
53
}
54

55
// DiffYAMLs compares two YAML documents and outputs the differences.
56
func DiffYAMLs(original, modified []byte) ([]byte, error) {
57
	var origNode, modNode yaml.Node
58
	if err := yaml.Unmarshal(original, &origNode); err != nil {
59
		return nil, err
60
	}
61
	if err := yaml.Unmarshal(modified, &modNode); err != nil {
62
		return nil, err
63
	}
64

65
	clearComments(&origNode)
66
	clearComments(&modNode)
67

68
	diff := compareNodes(origNode.Content[0], modNode.Content[0])
69
	if diff == nil {
70
		return []byte{}, nil
71
	}
72

73
	buffer := &bytes.Buffer{}
74
	encoder := yaml.NewEncoder(buffer)
75
	encoder.SetIndent(2)
76
	if err := encoder.Encode(diff); err != nil {
77
		return nil, err
78
	}
79
	encoder.Close()
80

81
	return buffer.Bytes(), nil
82
}
83

84
// clearComments cleans up comments in YAML nodes.
85
func clearComments(node *yaml.Node) {
86
	node.HeadComment = ""
87
	node.LineComment = ""
88
	node.FootComment = ""
89
	for _, n := range node.Content {
90
		clearComments(n)
91
	}
92
}
93

94
// compareNodes recursively finds differences between two YAML nodes.
95
func compareNodes(orig, mod *yaml.Node) *yaml.Node {
96
	if orig.Kind != mod.Kind {
97
		return mod
98
	}
99

100
	switch orig.Kind {
101
	case yaml.MappingNode:
102
		return compareMappingNodes(orig, mod)
103
	case yaml.SequenceNode:
104
		return compareSequenceNodes(orig, mod)
105
	case yaml.ScalarNode:
106
		if orig.Value != mod.Value {
107
			return mod
108
		}
109
	}
110
	return nil
111
}
112

113
// compareMappingNodes compares two mapping nodes and returns differences,
114
// prioritizing the order in the modified document but considering original document order where possible.
115
func compareMappingNodes(orig, mod *yaml.Node) *yaml.Node {
116
	diff := &yaml.Node{Kind: yaml.MappingNode}
117
	origMap := nodeMap(orig)
118
	modMap := nodeMap(mod)
119

120
	// Set to track keys from orig that have been processed
121
	processedKeys := make(map[string]bool)
122

123
	// First pass: iterate over keys in the modified node to maintain order
124
	for i := 0; i < len(mod.Content); i += 2 {
125
		key := mod.Content[i].Value
126
		modVal := modMap[key]
127
		origVal, origExists := origMap[key]
128

129
		if origExists {
130
			processedKeys[key] = true
131
			// Compare values for keys existing in both nodes
132
			changedNode := compareNodes(origVal, modVal)
133
			if changedNode != nil {
134
				addNodeToDiff(diff, key, changedNode)
135
			}
136
		} else {
137
			// New key in mod that doesn't exist in orig
138
			addNodeToDiff(diff, key, modVal)
139
		}
140
	}
141

142
	// Second pass: add keys from original that weren't in modified
143
	for i := 0; i < len(orig.Content); i += 2 {
144
		key := orig.Content[i].Value
145
		if !processedKeys[key] {
146
			origVal := origMap[key]
147
			addNodeToDiff(diff, key, origVal)
148
		}
149
	}
150

151
	if len(diff.Content) == 0 {
152
		return nil
153
	}
154
	return diff
155
}
156

157
// compareSequenceNodes compares two sequence nodes and returns differences.
158
func compareSequenceNodes(orig, mod *yaml.Node) *yaml.Node {
159
	diff := &yaml.Node{Kind: yaml.SequenceNode}
160
	origSet := nodeSet(orig)
161
	for _, modItem := range mod.Content {
162
		if !origSet[modItem.Value] {
163
			diff.Content = append(diff.Content, modItem)
164
		}
165
	}
166

167
	if len(diff.Content) == 0 {
168
		return nil
169
	}
170
	return diff
171
}
172

173
// nodeSet creates a set of values from sequence nodes.
174
func nodeSet(node *yaml.Node) map[string]bool {
175
	result := make(map[string]bool)
176
	for _, item := range node.Content {
177
		result[item.Value] = true
178
	}
179
	return result
180
}
181

182
// addNodeToDiff adds a node to the diff result.
183
func addNodeToDiff(diff *yaml.Node, key string, node *yaml.Node) {
184
	keyNode := &yaml.Node{Kind: yaml.ScalarNode, Value: key}
185
	diff.Content = append(diff.Content, keyNode)
186
	diff.Content = append(diff.Content, node)
187
}
188

189
// nodeMap creates a map from a YAML mapping node for easy lookup.
190
func nodeMap(node *yaml.Node) map[string]*yaml.Node {
191
	result := make(map[string]*yaml.Node)
192
	for i := 0; i+1 < len(node.Content); i += 2 {
193
		keyNode := node.Content[i]
194
		if keyNode.Kind == yaml.ScalarNode {
195
			result[keyNode.Value] = node.Content[i+1]
196
		}
197
	}
198
	return result
199
}
200