replacer.go

package csvheaders

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

// Replacer can replace headers. It embeds the provided csv.Reader
type Replacer struct {
	*csv.Reader

	inToOut map[string]string

	readOnce sync.Once
}

// NewReplacer builds a replacer on top of a csv reader.
// Usage: r := csvheaders.NewReplacer(csv.NewReader(in), headerReplacements)
func NewReplacer(reader *csv.Reader, inToOut map[string]string) *Replacer {
	return &Replacer{Reader: reader, inToOut: inToOut}
}

// NewReverseReplacer returns a replacer using mappings from struct definition to file contents.
// If there is a duplicate value in the headers mapping, this returns an error.
func NewReverseReplacer(reader *csv.Reader, outToIn map[string]string) (*Replacer, error) {
	r := &Replacer{Reader: reader, inToOut: make(map[string]string)}

	for out, in := range outToIn {
		if h, alreadyRegistered := r.inToOut[in]; alreadyRegistered {
			return nil, fmt.Errorf("struct tag %s registered with CSV headers %s and %s: %w", in, out, h, ErrDuplicateHeader)
		}

		r.inToOut[in] = out
	}

	return r, nil
}

// Read a record from the reader.
func (r *Replacer) Read() ([]string, error) {
	record, err := r.Reader.Read()
	if err != nil {
		return nil, fmt.Errorf("unable to Read from csv: %w", err)
	}

	// The headers appear on the first line - this Do is only executed on the first call to Read.
	r.readOnce.Do(func() {
		r.replaceHeaders(record)
	})

	return record, nil
}

// ReadAll wraps the CSV read all func.
func (r *Replacer) ReadAll() ([][]string, error) {
	records, err := r.Reader.ReadAll()
	if err != nil {
		return nil, fmt.Errorf("unable to ReadAll from csv: %w", err)
	}

	if len(records) == 0 {
		return records, nil
	}

	r.replaceHeaders(records[0])

	return records, nil
}

func (r *Replacer) replaceHeaders(headers []string) {
	// override with provided values.
	for idx, header := range headers {
		if newValue, found := r.inToOut[header]; found {
			headers[idx] = newValue
		}
	}
}