Changed to stream processing

noborus · Mar 14, 2023 · 442522b · 442522b
1 parent 60e819a
commit 442522b
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 33 deletions.
diff --git a/cmd/csv.go b/cmd/csv.go
@@ -31,11 +31,14 @@ func toCSV(delimiter rune) {
 	g.Header = header - 1
 	g.LimitSplit = LimitSplit
 	g.TrimSpace = true
-	table := g.Rows()
 
 	w := csv.NewWriter(os.Stdout)
 	w.Comma = delimiter
-	for _, record := range table {
+	for {
+		record, err := g.Read()
+		if err != nil {
+			break
+		}
 		if err := w.Write(record); err != nil {
 			log.Fatal(err)
 		}

diff --git a/cmd/root.go b/cmd/root.go
@@ -30,13 +30,15 @@ func writeTable(args []string) {
 	g.Header = header - 1
 	g.LimitSplit = LimitSplit
 	g.TrimSpace = false
-	table := g.Rows()
-
-	write(table)
+	write(g)
 }
 
-func write(table [][]string) {
-	for _, row := range table {
+func write(g *guesswidth.GuessWidth) {
+	for {
+		row, err := g.Read()
+		if err != nil {
+			break
+		}
 		for n, col := range row {
 			if n > 0 {
 				fmt.Print(fence)

diff --git a/guesswidth.go b/guesswidth.go
@@ -10,59 +10,83 @@ import (
 	"github.com/mattn/go-runewidth"
 )
 
-type guessWidth struct {
+type GuessWidth struct {
 	scanner    *bufio.Scanner
-	preLines   []string
+	pos        []int
+	preRows    [][]string
 	preNum     int
+	preCount   int
 	Header     int
 	LimitSplit int
 	TrimSpace  bool
 }
 
-var preNum = 10
+var PreNum = 10
 
-func New(r io.Reader) *guessWidth {
+func New(r io.Reader) *GuessWidth {
 	scanner := bufio.NewScanner(r)
-	lines := preRead(scanner, preNum)
-	return &guessWidth{
+	g := &GuessWidth{
 		scanner:    scanner,
-		preLines:   lines,
-		preNum:     preNum,
+		preNum:     PreNum,
 		Header:     0,
 		LimitSplit: 0,
 		TrimSpace:  true,
 	}
+	return g
 }
 
-func (g *guessWidth) Rows() [][]string {
-	pos := widthPositions(g.preLines, g.Header)
+func (g *GuessWidth) ReadAll() [][]string {
+	g.Scan(g.preNum)
+
 	var rows [][]string
-	if g.LimitSplit > 0 {
-		if len(pos) > g.LimitSplit {
-			pos = pos[:g.LimitSplit]
+	for {
+		columns, err := g.Read()
+		if err != nil {
+			break
 		}
-		rows = toRows(g.preLines, pos, g.TrimSpace)
-	} else {
-		rows = toRows(g.preLines, pos, g.TrimSpace)
-	}
-
-	for g.scanner.Scan() {
-		line := g.scanner.Text()
-		columns := split(line, pos, g.TrimSpace)
 		rows = append(rows, columns)
 	}
 	return rows
 }
 
-func preRead(scanner *bufio.Scanner, preRead int) []string {
+func (g *GuessWidth) Scan(preRead int) {
 	var lines []string
 	for i := 0; i < preRead; i++ {
-		if !scanner.Scan() {
-			return lines
+		if !g.scanner.Scan() {
+			break
 		}
-		lines = append(lines, scanner.Text())
+		lines = append(lines, g.scanner.Text())
+	}
+
+	g.pos = widthPositions(lines, g.Header)
+
+	var rows [][]string
+	if g.LimitSplit > 0 {
+		if len(g.pos) > g.LimitSplit {
+			g.pos = g.pos[:g.LimitSplit]
+		}
+		rows = toRows(lines, g.pos, g.TrimSpace)
+	} else {
+		rows = toRows(lines, g.pos, g.TrimSpace)
+	}
+	g.preRows = rows
+	return
+}
+
+func (g *GuessWidth) Read() ([]string, error) {
+	if len(g.preRows) == 0 {
+		g.Scan(PreNum)
+	}
+	if g.preCount < len(g.preRows) {
+		row := g.preRows[g.preCount]
+		g.preCount++
+		return row, nil
+	}
+	if !g.scanner.Scan() {
+		return nil, fmt.Errorf("EOF")
 	}
-	return lines
+	line := g.scanner.Text()
+	return split(line, g.pos, g.TrimSpace), nil
 }
 
 func ToTable(lines []string, header int, trimSpace bool) [][]string {