lazygit/pkg/utils/yaml_utils/yaml_utils.go

package yaml_utils

import (
	"errors"
	"fmt"

	"gopkg.in/yaml.v3"
)

// takes a yaml document in bytes, a path to a key, and a value to set. The value must be a scalar.
func UpdateYamlValue(yamlBytes []byte, path []string, value string) ([]byte, error) {
	// Parse the YAML file.
	var node yaml.Node
	err := yaml.Unmarshal(yamlBytes, &node)
	if err != nil {
		return nil, fmt.Errorf("failed to parse YAML: %w", err)
	}

	// Empty document: need to create the top-level map ourselves
	if len(node.Content) == 0 {
		node.Content = append(node.Content, &yaml.Node{
			Kind: yaml.MappingNode,
		})
	}

	body := node.Content[0]

	if body.Kind != yaml.MappingNode {
		return yamlBytes, errors.New("yaml document is not a dictionary")
	}

	if err := updateYamlNode(body, path, value); err != nil {
		return yamlBytes, err
	}

	// Convert the updated YAML node back to YAML bytes.
	updatedYAMLBytes, err := yaml.Marshal(body)
	if err != nil {
		return nil, fmt.Errorf("failed to convert YAML node to bytes: %w", err)
	}

	return updatedYAMLBytes, nil
}

// Recursive function to update the YAML node.
func updateYamlNode(node *yaml.Node, path []string, value string) error {
	if len(path) == 0 {
		if node.Kind != yaml.ScalarNode {
			return errors.New("yaml node is not a scalar")
		}
		node.Value = value
		return nil
	}

	if node.Kind != yaml.MappingNode {
		return errors.New("yaml node in path is not a dictionary")
	}

	key := path[0]
	for i := 0; i < len(node.Content)-1; i += 2 {
		if node.Content[i].Value == key {
			return updateYamlNode(node.Content[i+1], path[1:], value)
		}
	}

	// if the key doesn't exist, we'll add it

	// at end of path: add the new key, done
	if len(path) == 1 {
		node.Content = append(node.Content, &yaml.Node{
			Kind:  yaml.ScalarNode,
			Value: key,
		}, &yaml.Node{
			Kind:  yaml.ScalarNode,
			Value: value,
		})
		return nil
	}

	// otherwise, create the missing intermediate node and continue
	newNode := &yaml.Node{
		Kind: yaml.MappingNode,
	}
	node.Content = append(node.Content, &yaml.Node{
		Kind:  yaml.ScalarNode,
		Value: key,
	}, newNode)
	return updateYamlNode(newNode, path[1:], value)
}