miller/pkg/output/record_writer_pprint.go

package output

import (
	"bufio"
	"strings"

	"github.com/johnkerl/miller/v6/pkg/cli"
	"github.com/johnkerl/miller/v6/pkg/colorizer"
	"github.com/johnkerl/miller/v6/pkg/lib"
	"github.com/johnkerl/miller/v6/pkg/mlrval"
	"github.com/johnkerl/miller/v6/pkg/types"
)

type RecordWriterPPRINT struct {
	writerOptions *cli.TWriterOptions
	// Input:
	records []*mlrval.Mlrmap

	// State:
	lastJoinedHeader *string
	batch            []*mlrval.Mlrmap
}

func NewRecordWriterPPRINT(writerOptions *cli.TWriterOptions) (*RecordWriterPPRINT, error) {
	return &RecordWriterPPRINT{
		writerOptions: writerOptions,
		records:       []*mlrval.Mlrmap{},

		lastJoinedHeader: nil,
		batch:            []*mlrval.Mlrmap{},
	}, nil
}

func (writer *RecordWriterPPRINT) Write(
	outrec *mlrval.Mlrmap,
	_ *types.Context,
	bufferedOutputStream *bufio.Writer,
	outputIsStdout bool,
) error {
	// Group records by have-same-schema or not. Pretty-print each
	// homoegeneous sublist, or "batch".
	//
	// No output until end of a homogeneous batch of records, since we need to
	// find out max width down each column.

	if outrec != nil { // Not end of record stream
		if writer.lastJoinedHeader == nil {
			// First output record:
			// * New batch
			// * No old batch to print
			writer.batch = append(writer.batch, outrec)
			temp := strings.Join(outrec.GetKeys(), ",")
			writer.lastJoinedHeader = &temp
		} else {
			// May or may not continue the same homogeneous batch
			joinedHeader := strings.Join(outrec.GetKeys(), ",")
			if *writer.lastJoinedHeader != joinedHeader {
				// Print and free old batch
				nonEmpty := writer.writeHeterogenousList(
					writer.batch,
					writer.writerOptions.BarredPprintOutput,
					bufferedOutputStream,
					outputIsStdout,
				)
				if nonEmpty {
					// Print a newline
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
				// Start a new batch
				writer.batch = []*mlrval.Mlrmap{outrec}
				writer.lastJoinedHeader = &joinedHeader
			} else {
				// Continue the batch
				writer.batch = append(writer.batch, outrec)
			}
		}

	} else { // End of record stream
		if len(writer.batch) > 0 {
			writer.writeHeterogenousList(writer.batch, writer.writerOptions.BarredPprintOutput,
				bufferedOutputStream, outputIsStdout)
		}
	}

	return nil
}

// Returns false if there was nothing but empty record(s), e.g. 'mlr gap -n 10'.
func (writer *RecordWriterPPRINT) writeHeterogenousList(
	records []*mlrval.Mlrmap,
	barred bool,
	bufferedOutputStream *bufio.Writer,
	outputIsStdout bool,
) bool {
	maxWidths := make(map[string]int)
	var maxNR int64 = 0

	for _, outrec := range records {
		nr := outrec.FieldCount
		if maxNR < nr {
			maxNR = nr
		}
		for pe := outrec.Head; pe != nil; pe = pe.Next {
			width := lib.DisplayWidth(pe.Value.String())
			if width == 0 {
				width = 1 // We'll rewrite "" to "-" below
			}
			oldMaxWidth := maxWidths[pe.Key]
			if width > oldMaxWidth {
				maxWidths[pe.Key] = width
			}
		}
	}

	if maxNR == 0 {
		return false
	}
	// Column name may be longer/shorter than all data values in the column
	for key, oldMaxWidth := range maxWidths {
		width := lib.DisplayWidth(key)
		if width > oldMaxWidth {
			maxWidths[key] = width
		}
	}
	if barred {
		writer.writeHeterogenousListBarred(records, maxWidths, bufferedOutputStream, outputIsStdout)
	} else {
		writer.writeHeterogenousListNonBarred(records, maxWidths, bufferedOutputStream, outputIsStdout)
	}
	return true
}

// Example:
//
// a   b   i  x                    y
// pan pan 1  0.3467901443380824   0.7268028627434533
// eks pan 2  -0.7586799647899636  0.5221511083334797
// wye wye 3  0.20460330576630303  0.33831852551664776
// eks wye 4  -0.38139939387114097 0.13418874328430463
// wye pan 5  0.5732889198020006   0.8636244699032729

func (writer *RecordWriterPPRINT) writeHeterogenousListNonBarred(
	records []*mlrval.Mlrmap,
	maxWidths map[string]int,
	bufferedOutputStream *bufio.Writer,
	outputIsStdout bool,
) {

	onFirst := true
	for _, outrec := range records {

		// Print header line
		if onFirst && !writer.writerOptions.HeaderlessOutput {
			for pe := outrec.Head; pe != nil; pe = pe.Next {
				if !writer.writerOptions.RightAlignedPPRINTOutput { // left-align
					if pe.Next != nil {
						// Header line, left-align, not last column
						bufferedOutputStream.WriteString(colorizer.MaybeColorizeKey(pe.Key, outputIsStdout))
						writer.writePadding(pe.Key, maxWidths[pe.Key], bufferedOutputStream)
						bufferedOutputStream.WriteString(writer.writerOptions.OFS)
					} else {
						// Header line, left-align, last column
						bufferedOutputStream.WriteString(colorizer.MaybeColorizeKey(pe.Key, outputIsStdout))
						bufferedOutputStream.WriteString(writer.writerOptions.ORS)
					}
				} else { // right-align
					writer.writePadding(pe.Key, maxWidths[pe.Key], bufferedOutputStream)
					bufferedOutputStream.WriteString(colorizer.MaybeColorizeKey(pe.Key, outputIsStdout))
					if pe.Next != nil {
						// Header line, right-align, not last column
						bufferedOutputStream.WriteString(writer.writerOptions.OFS)
					} else {
						// Header line, right-align, last column
						bufferedOutputStream.WriteString(writer.writerOptions.ORS)
					}
				}

			}
		}
		onFirst = false

		// Print data lines
		for pe := outrec.Head; pe != nil; pe = pe.Next {
			s := pe.Value.String()
			if s == "" {
				s = "-"
			}
			if !writer.writerOptions.RightAlignedPPRINTOutput { // left-align
				if pe.Next != nil {
					// Data line, left-align, not last column
					bufferedOutputStream.WriteString(colorizer.MaybeColorizeValue(s, outputIsStdout))
					writer.writePadding(s, maxWidths[pe.Key], bufferedOutputStream)
					bufferedOutputStream.WriteString(writer.writerOptions.OFS)
				} else {
					// Data line, left-align, last column
					bufferedOutputStream.WriteString(colorizer.MaybeColorizeValue(s, outputIsStdout))
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			} else { // right-align
				writer.writePadding(s, maxWidths[pe.Key], bufferedOutputStream)
				bufferedOutputStream.WriteString(colorizer.MaybeColorizeValue(s, outputIsStdout))
				if pe.Next != nil {
					// Data line, right-align, not last column
					bufferedOutputStream.WriteString(writer.writerOptions.OFS)
				} else {
					// Data line, right-align, last column
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			}
		}

		if writer.writerOptions.FlushOnEveryRecord {
			// bufio.Writer errors are sticky; the final Flush in pkg/stream is checked
			_ = bufferedOutputStream.Flush()
		}
	}
}

// Example:
//
// +-----+-----+----+----------------------+---------------------+
// | a   | b   | i  | x                    | y                   |
// +-----+-----+----+----------------------+---------------------+
// | pan | pan | 1  | 0.3467901443380824   | 0.7268028627434533  |
// | eks | pan | 2  | -0.7586799647899636  | 0.5221511083334797  |
// | wye | wye | 3  | 0.20460330576630303  | 0.33831852551664776 |
// | eks | wye | 4  | -0.38139939387114097 | 0.13418874328430463 |
// | wye | pan | 5  | 0.5732889198020006   | 0.8636244699032729  |
// +-----+-----+----+----------------------+---------------------+

type barredChars struct {
	horizontalStart          string
	horizontalMiddle         string
	horizontalEnd            string
	horizontalBar            string
	verticalStart            string
	verticalMiddle           string
	verticalEnd              string
	firstRowHorizontalStart  string
	firstRowHorizontalMiddle string
	firstRowHorizontalEnd    string
	lastRowHorizontalStart   string
	lastRowHorizontalMiddle  string
	lastRowHorizontalEnd     string
}

func (writer *RecordWriterPPRINT) writeHeterogenousListBarred(
	records []*mlrval.Mlrmap,
	maxWidths map[string]int,
	bufferedOutputStream *bufio.Writer,
	outputIsStdout bool,
) {

	bc := func() barredChars {
		ofs := writer.writerOptions.OFS
		if writer.writerOptions.BarredUseUnicode {
			return barredChars{
				horizontalStart:          "├─",
				horizontalMiddle:         "─┼─",
				horizontalEnd:            "─┤",
				horizontalBar:            "─",
				verticalStart:            "│" + ofs,
				verticalMiddle:           ofs + "│" + ofs,
				verticalEnd:              ofs + "│",
				firstRowHorizontalStart:  "┌─",
				firstRowHorizontalMiddle: "─┬─",
				firstRowHorizontalEnd:    "─┐",
				lastRowHorizontalStart:   "└─",
				lastRowHorizontalMiddle:  "─┴─",
				lastRowHorizontalEnd:     "─┘",
			}
		} else {
			return barredChars{
				horizontalStart:          "+-",
				horizontalMiddle:         "-+-",
				horizontalEnd:            "-+",
				horizontalBar:            "-",
				verticalStart:            "|" + ofs,
				verticalMiddle:           ofs + "|" + ofs,
				verticalEnd:              ofs + "|",
				firstRowHorizontalStart:  "+-",
				firstRowHorizontalMiddle: "-+-",
				firstRowHorizontalEnd:    "-+",
				lastRowHorizontalStart:   "+-",
				lastRowHorizontalMiddle:  "-+-",
				lastRowHorizontalEnd:     "-+",
			}
		}
	}()

	horizontalBars := make(map[string]string)
	for key, width := range maxWidths {
		horizontalBars[key] = strings.Repeat(bc.horizontalBar, width)
	}

	onFirst := true
	for i, outrec := range records {

		// Print header line
		if onFirst && !writer.writerOptions.HeaderlessOutput {
			bufferedOutputStream.WriteString(bc.firstRowHorizontalStart)
			for pe := outrec.Head; pe != nil; pe = pe.Next {
				bufferedOutputStream.WriteString(horizontalBars[pe.Key])
				if pe.Next != nil {
					bufferedOutputStream.WriteString(bc.firstRowHorizontalMiddle)
				} else {
					bufferedOutputStream.WriteString(bc.firstRowHorizontalEnd)
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			}

			bufferedOutputStream.WriteString(bc.verticalStart)
			for pe := outrec.Head; pe != nil; pe = pe.Next {
				if !writer.writerOptions.RightAlignedPPRINTOutput { // left-align
					bufferedOutputStream.WriteString(colorizer.MaybeColorizeKey(pe.Key, outputIsStdout))
					writer.writePadding(pe.Key, maxWidths[pe.Key], bufferedOutputStream)
				} else { // right-align
					writer.writePadding(pe.Key, maxWidths[pe.Key], bufferedOutputStream)
					bufferedOutputStream.WriteString(colorizer.MaybeColorizeKey(pe.Key, outputIsStdout))
				}
				if pe.Next != nil {
					bufferedOutputStream.WriteString(bc.verticalMiddle)
				} else {
					bufferedOutputStream.WriteString(bc.verticalEnd)
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			}

			bufferedOutputStream.WriteString(bc.horizontalStart)
			for pe := outrec.Head; pe != nil; pe = pe.Next {
				bufferedOutputStream.WriteString(horizontalBars[pe.Key])
				if pe.Next != nil {
					bufferedOutputStream.WriteString(bc.horizontalMiddle)
				} else {
					bufferedOutputStream.WriteString(bc.horizontalEnd)
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			}
		}
		onFirst = false

		// Print data lines
		bufferedOutputStream.WriteString(bc.verticalStart)
		for pe := outrec.Head; pe != nil; pe = pe.Next {
			s := pe.Value.String()
			if !writer.writerOptions.RightAlignedPPRINTOutput { // left-align
				bufferedOutputStream.WriteString(colorizer.MaybeColorizeValue(s, outputIsStdout))
				writer.writePadding(s, maxWidths[pe.Key], bufferedOutputStream)
			} else { // right-align
				writer.writePadding(s, maxWidths[pe.Key], bufferedOutputStream)
				bufferedOutputStream.WriteString(colorizer.MaybeColorizeValue(s, outputIsStdout))
			}
			if pe.Next != nil {
				bufferedOutputStream.WriteString(bc.verticalMiddle)
			} else {
				bufferedOutputStream.WriteString(bc.verticalEnd)
				bufferedOutputStream.WriteString(writer.writerOptions.ORS)
			}
		}

		// Print last line
		if i == len(records)-1 {
			bufferedOutputStream.WriteString(bc.lastRowHorizontalStart)
			for pe := outrec.Head; pe != nil; pe = pe.Next {
				bufferedOutputStream.WriteString(horizontalBars[pe.Key])
				if pe.Next != nil {
					bufferedOutputStream.WriteString(bc.lastRowHorizontalMiddle)
				} else {
					bufferedOutputStream.WriteString(bc.lastRowHorizontalEnd)
					bufferedOutputStream.WriteString(writer.writerOptions.ORS)
				}
			}
		}

		if writer.writerOptions.FlushOnEveryRecord {
			// bufio.Writer errors are sticky; the final Flush in pkg/stream is checked
			_ = bufferedOutputStream.Flush()
		}
	}
}

func (writer *RecordWriterPPRINT) writePadding(
	text string,
	fieldWidth int,
	bufferedOutputStream *bufio.Writer,
) {
	textWidth := lib.DisplayWidth(text)
	padWidth := fieldWidth - textWidth
	if padWidth > 0 {
		bufferedOutputStream.WriteString(strings.Repeat(" ", padWidth))
	}
}