yq/pkg/yqlib/encoder_csv.go

package yqlib

import (
	"encoding/csv"
	"fmt"
	"io"
)

type csvEncoder struct {
	separator rune
}

func NewCsvEncoder(separator rune) Encoder {
	return &csvEncoder{separator: separator}
}

func (e *csvEncoder) CanHandleAliases() bool {
	return false
}

func (e *csvEncoder) PrintDocumentSeparator(writer io.Writer) error {
	return nil
}

func (e *csvEncoder) PrintLeadingContent(writer io.Writer, content string) error {
	return nil
}

func (e *csvEncoder) encodeRow(csvWriter *csv.Writer, contents []*CandidateNode) error {
	stringValues := make([]string, len(contents))

	for i, child := range contents {

		if child.Kind != ScalarNode {
			return fmt.Errorf("csv encoding only works for arrays of scalars (string/numbers/booleans), child[%v] is a %v", i, child.Tag)
		}
		stringValues[i] = child.Value
	}
	return csvWriter.Write(stringValues)
}

func (e *csvEncoder) encodeArrays(csvWriter *csv.Writer, content []*CandidateNode) error {
	for i, child := range content {

		if child.Kind != SequenceNode {
			return fmt.Errorf("csv encoding only works for arrays of scalars (string/numbers/booleans), child[%v] is a %v", i, child.Tag)
		}
		err := e.encodeRow(csvWriter, child.Content)
		if err != nil {
			return err
		}
	}
	return nil
}

func (e *csvEncoder) extractHeader(child *CandidateNode) ([]*CandidateNode, error) {
	if child.Kind != MappingNode {
		return nil, fmt.Errorf("csv object encoding only works for arrays of flat objects (string key => string/numbers/boolean value), child[0] is a %v", child.Tag)
	}
	mapKeys := getMapKeys(child)
	return mapKeys.Content, nil
}

func (e *csvEncoder) createChildRow(child *CandidateNode, headers []*CandidateNode) []*CandidateNode {
	childRow := make([]*CandidateNode, 0)
	for _, header := range headers {
		keyIndex := findKeyInMap(child, header)
		value := createScalarNode(nil, "")
		if keyIndex != -1 {
			value = child.Content[keyIndex+1]
		}
		childRow = append(childRow, value)
	}
	return childRow

}

func (e *csvEncoder) encodeObjects(csvWriter *csv.Writer, content []*CandidateNode) error {
	headers, err := e.extractHeader(content[0])
	if err != nil {
		return nil
	}

	err = e.encodeRow(csvWriter, headers)
	if err != nil {
		return nil
	}

	for i, child := range content {
		if child.Kind != MappingNode {
			return fmt.Errorf("csv object encoding only works for arrays of flat objects (string key => string/numbers/boolean value), child[%v] is a %v", i, child.Tag)
		}
		row := e.createChildRow(child, headers)
		err = e.encodeRow(csvWriter, row)
		if err != nil {
			return err
		}

	}
	return nil
}

func (e *csvEncoder) Encode(writer io.Writer, originalNode *CandidateNode) error {
	if originalNode.Kind == ScalarNode {
		return writeString(writer, originalNode.Value+"\n")
	}

	csvWriter := csv.NewWriter(writer)
	csvWriter.Comma = e.separator

	// node must be a sequence
	node := originalNode.unwrapDocument()
	if node.Kind != SequenceNode {
		return fmt.Errorf("csv encoding only works for arrays, got: %v", node.Tag)
	} else if len(node.Content) == 0 {
		return nil
	}
	if node.Content[0].Kind == ScalarNode {
		return e.encodeRow(csvWriter, node.Content)
	}

	if node.Content[0].Kind == MappingNode {
		return e.encodeObjects(csvWriter, node.Content)
	}

	return e.encodeArrays(csvWriter, node.Content)

}