mirror of
https://github.com/johnkerl/miller.git
synced 2026-01-23 02:14:13 +00:00
100 lines
2.7 KiB
Go
100 lines
2.7 KiB
Go
package output
|
|
|
|
import (
|
|
"bufio"
|
|
"fmt"
|
|
"strings"
|
|
|
|
csv "github.com/johnkerl/miller/v6/pkg/go-csv"
|
|
|
|
"github.com/johnkerl/miller/v6/pkg/cli"
|
|
"github.com/johnkerl/miller/v6/pkg/mlrval"
|
|
"github.com/johnkerl/miller/v6/pkg/types"
|
|
)
|
|
|
|
type RecordWriterCSV struct {
|
|
writerOptions *cli.TWriterOptions
|
|
csvWriter *csv.Writer
|
|
needToPrintHeader bool
|
|
firstRecordKeys []string
|
|
firstRecordNF int64
|
|
quoteAll bool // For double-quote around all fields
|
|
}
|
|
|
|
func NewRecordWriterCSV(writerOptions *cli.TWriterOptions) (*RecordWriterCSV, error) {
|
|
if len(writerOptions.OFS) != 1 {
|
|
return nil, fmt.Errorf("for CSV, OFS can only be a single character")
|
|
}
|
|
if writerOptions.ORS != "\n" && writerOptions.ORS != "\r\n" {
|
|
return nil, fmt.Errorf("for CSV, ORS cannot be altered")
|
|
}
|
|
writer := &RecordWriterCSV{
|
|
writerOptions: writerOptions,
|
|
csvWriter: nil, // will be set on first Write() wherein we have the output stream
|
|
needToPrintHeader: !writerOptions.HeaderlessOutput,
|
|
firstRecordKeys: nil,
|
|
firstRecordNF: -1,
|
|
quoteAll: writerOptions.CSVQuoteAll,
|
|
}
|
|
return writer, nil
|
|
}
|
|
|
|
func (writer *RecordWriterCSV) Write(
|
|
outrec *mlrval.Mlrmap,
|
|
_ *types.Context,
|
|
bufferedOutputStream *bufio.Writer,
|
|
outputIsStdout bool,
|
|
) error {
|
|
if outrec == nil {
|
|
// End of record stream: nothing special for this output format
|
|
return nil
|
|
}
|
|
|
|
if writer.csvWriter == nil {
|
|
writer.csvWriter = csv.NewWriter(bufferedOutputStream)
|
|
writer.csvWriter.Comma = rune(writer.writerOptions.OFS[0]) // xxx temp
|
|
}
|
|
|
|
if writer.firstRecordKeys == nil {
|
|
writer.firstRecordKeys = outrec.GetKeys()
|
|
writer.firstRecordNF = int64(len(writer.firstRecordKeys))
|
|
}
|
|
|
|
if writer.needToPrintHeader {
|
|
fields := make([]string, outrec.FieldCount)
|
|
i := 0
|
|
for pe := outrec.Head; pe != nil; pe = pe.Next {
|
|
fields[i] = pe.Key
|
|
i++
|
|
}
|
|
writer.WriteCSVRecordMaybeColorized(fields, bufferedOutputStream, outputIsStdout, true, writer.quoteAll)
|
|
writer.needToPrintHeader = false
|
|
}
|
|
|
|
var outputNF int64 = outrec.FieldCount
|
|
if outputNF < writer.firstRecordNF {
|
|
outputNF = writer.firstRecordNF
|
|
}
|
|
|
|
fields := make([]string, outputNF)
|
|
var i int64 = 0
|
|
for pe := outrec.Head; pe != nil; pe = pe.Next {
|
|
if i < writer.firstRecordNF && pe.Key != writer.firstRecordKeys[i] {
|
|
return fmt.Errorf(
|
|
"CSV schema change: first keys \"%s\"; current keys \"%s\"",
|
|
strings.Join(writer.firstRecordKeys, writer.writerOptions.OFS),
|
|
strings.Join(outrec.GetKeys(), writer.writerOptions.OFS),
|
|
)
|
|
}
|
|
fields[i] = pe.Value.String()
|
|
i++
|
|
}
|
|
|
|
for ; i < outputNF; i++ {
|
|
fields[i] = ""
|
|
}
|
|
|
|
writer.WriteCSVRecordMaybeColorized(fields, bufferedOutputStream, outputIsStdout, false, writer.quoteAll)
|
|
|
|
return nil
|
|
}
|