github.com/verygoodsoftwarenotvirus/cartogopher@v0.0.0-20160213112503-f5fe2b6d4bd1/writer.go (about)

     1  package cartogopher
     2  
     3  import (
     4  	"encoding/csv"
     5  	"fmt"
     6  	"io"
     7  )
     8  
     9  // MapWriter mimics the writer struct.
    10  type MapWriter struct {
    11  	InputHeaders    []string
    12  	OutputHeaderMap map[string]int
    13  	Writer          *csv.Writer
    14  }
    15  
    16  func (w *MapWriter) createOutputHeaderMap() {
    17  	result := map[string]int{}
    18  	for index, header := range w.InputHeaders {
    19  		result[header] = index
    20  	}
    21  	w.OutputHeaderMap = result
    22  }
    23  
    24  func (w MapWriter) createOutputSlice(row map[string]string) ([]string, error) {
    25  	outputSlice := make([]string, len(w.OutputHeaderMap))
    26  	for header, value := range row {
    27  		if _, ok := w.OutputHeaderMap[header]; ok {
    28  			outputSlice[w.OutputHeaderMap[header]] = value
    29  		} else {
    30  			return []string{}, fmt.Errorf("Provided row contains invalid header field: %v", header)
    31  		}
    32  	}
    33  	return outputSlice, nil
    34  }
    35  
    36  // Write recreates the built-in CSV writer's Write method
    37  func (w MapWriter) Write(row map[string]string) error {
    38  	if len(row) > len(w.OutputHeaderMap) {
    39  		return fmt.Errorf("Provided row has %v fields, whereas there are only %v headers", len(row), len(w.OutputHeaderMap))
    40  	}
    41  	outputSlice, err := w.createOutputSlice(row)
    42  	if err != nil {
    43  		return err
    44  	}
    45  	err = w.Writer.Write(outputSlice)
    46  	return err
    47  }
    48  
    49  // WriteAll rereates the built-in CSV writer's WriteAll method
    50  func (w MapWriter) WriteAll(rows []map[string]string) error {
    51  	for _, row := range rows {
    52  		err := w.Write(row)
    53  		if err != nil {
    54  			return err
    55  		}
    56  	}
    57  	return nil
    58  }
    59  
    60  // Flush simply calls the built-in CSV writer's flush method
    61  func (w *MapWriter) Flush() {
    62  	w.Writer.Flush()
    63  }
    64  
    65  // NewWriter creates a new writer
    66  func NewWriter(w io.Writer, headers []string) *MapWriter {
    67  	writer := csv.NewWriter(w)
    68  	result := &MapWriter{
    69  		InputHeaders: headers,
    70  		Writer:       writer,
    71  	}
    72  	result.Writer.Write(result.InputHeaders)
    73  	result.createOutputHeaderMap()
    74  	return result
    75  }