Merge pull request #295 from olekukonko/counter

Counter Implementation

Author: olekukonko

README.md

@@ -28,7 +28,7 @@ go get github.com/olekukonko/[email protected]
 #### Latest  Version
 The latest stable version
 ```bash
-go get github.com/olekukonko/tablewriter@v1.0.9
+go get github.com/olekukonko/tablewriter@v1.1.0
 ```
 
 **Warning:** Version `v1.0.0` contains missing functionality and should not be used.
@@ -62,7 +62,7 @@ func main() {
 	data := [][]string{
 		{"Package", "Version", "Status"},
 		{"tablewriter", "v0.0.5", "legacy"},
-		{"tablewriter", "v1.0.9", "latest"},
+		{"tablewriter", "v1.1.0", "latest"},
 	}
 
 	table := tablewriter.NewWriter(os.Stdout)
@@ -77,7 +77,7 @@ func main() {
 │   PACKAGE   │ VERSION │ STATUS │
 ├─────────────┼─────────┼────────┤
 │ tablewriter │ v0.0.5  │ legacy │
-│ tablewriter │ v1.0.9  │ latest │
+│ tablewriter │ v1.1.0  │ latest │
 └─────────────┴─────────┴────────┘
 ```
 
@@ -1081,6 +1081,8 @@ func (t Time) Format() string {
 
 - `AutoFormat` changes See [#261](https://github.com/olekukonko/tablewriter/issues/261)
 
+## What is new
+- `Counting` changes See [#294](https://github.com/olekukonko/tablewriter/issues/294)
 
 ## Command-Line Tool
 

config.go

@@ -14,6 +14,7 @@ type Config struct {
 	Stream   tw.StreamConfig
 	Behavior tw.Behavior
 	Widths   tw.CellWidth
+	Counter  tw.Counter
 }
 
 // ConfigBuilder provides a fluent interface for building Config

option.go

@@ -671,6 +671,31 @@ func WithSymbols(symbols tw.Symbols) Option {
 	}
 }
 
+// WithCounters enables line counting by wrapping the table's writer.
+// If a custom counter (that implements tw.Counter) is provided, it will be used.
+// If the provided counter is nil, a default tw.LineCounter will be used.
+// The final count can be retrieved via the table.Lines() method after Render() is called.
+func WithCounters(counters ...tw.Counter) Option {
+	return func(target *Table) {
+		// Iterate through the provided counters and add any non-nil ones.
+		for _, c := range counters {
+			if c != nil {
+				target.counters = append(target.counters, c)
+			}
+		}
+	}
+}
+
+// WithLineCounter enables the default line counter.
+// A new instance of tw.LineCounter is added to the table's list of counters.
+// The total count can be retrieved via the table.Lines() method after Render() is called.
+func WithLineCounter() Option {
+	return func(target *Table) {
+		// Important: Create a new instance so tables don't share counters.
+		target.counters = append(target.counters, &tw.LineCounter{})
+	}
+}
+
 // defaultConfig returns a default Config with sensible settings for headers, rows, footers, and behavior.
 func defaultConfig() Config {
 	return Config{

tablewriter.go

@@ -22,6 +22,7 @@ import (
 // Table represents a table instance with content and rendering capabilities.
 type Table struct {
 	writer       io.Writer           // Destination for table output
+	counters     []tw.Counter        // Counters for indices
 	rows         [][][]string        // Row data, supporting multi-line cells
 	headers      [][]string          // Header content
 	footers      [][]string          // Footer content
@@ -510,6 +511,28 @@ func (t *Table) Render() error {
 	return t.render()
 }
 
+// Lines returns the total number of lines rendered.
+// This method is only effective if the WithLineCounter() option was used during
+// table initialization and must be called *after* Render().
+// It actively searches for the default tw.LineCounter among all active counters.
+// It returns -1 if the line counter was not enabled.
+func (t *Table) Lines() int {
+	for _, counter := range t.counters {
+		if lc, ok := counter.(*tw.LineCounter); ok {
+			return lc.Total()
+		}
+	}
+	// use -1 to indicate no line counter is attached
+	return -1
+}
+
+// Counters returns the slice of all active counter instances.
+// This is useful when multiple counters are enabled.
+// It must be called *after* Render().
+func (t *Table) Counters() []tw.Counter {
+	return t.counters
+}
+
 // Trimmer trims whitespace from a string based on the Table’s configuration.
 // It conditionally applies strings.TrimSpace to the input string if the TrimSpace behavior
 // is enabled in t.config.Behavior, otherwise returning the string unchanged. This method
@@ -1361,20 +1384,38 @@ func (t *Table) prepareWithMerges(content [][]string, config tw.CellConfig, posi
 func (t *Table) render() error {
 	t.ensureInitialized()
 
+	// Save the original writer and schedule its restoration upon function exit.
+	// This guarantees the table's writer is restored even if errors occur.
+	originalWriter := t.writer
+	defer func() {
+		t.writer = originalWriter
+	}()
+
+	// If a counter is active, wrap the writer in a MultiWriter.
+	if len(t.counters) > 0 {
+		// The slice must be of type io.Writer.
+		// Start it with the original destination writer.
+		allWriters := []io.Writer{originalWriter}
+
+		// Append each counter to the slice of writers.
+		for _, c := range t.counters {
+			allWriters = append(allWriters, c)
+		}
+
+		// Create a MultiWriter that broadcasts to the original writer AND all counters.
+		t.writer = io.MultiWriter(allWriters...)
+	}
+
 	if t.config.Stream.Enable {
 		t.logger.Warn("Render() called in streaming mode. Use Start/Append/Close methods instead.")
 		return errors.New("render called in streaming mode; use Start/Append/Close")
 	}
 
-	// Calculate and cache numCols for THIS batch render pass
-	t.batchRenderNumCols = t.maxColumns() // Calculate ONCE
-	t.isBatchRenderNumColsSet = true      // Mark the cache as active for this render pass
-	t.logger.Debugf("Render(): Set batchRenderNumCols to %d and isBatchRenderNumColsSet to true.", t.batchRenderNumCols)
-
+	// Calculate and cache the column count for this specific batch render pass.
+	t.batchRenderNumCols = t.maxColumns()
+	t.isBatchRenderNumColsSet = true
 	defer func() {
 		t.isBatchRenderNumColsSet = false
-		// t.batchRenderNumCols = 0; // Optional: reset to 0, or leave as is.
-		// Since isBatchRenderNumColsSet is false, its value won't be used by getNumColsToUse.
 		t.logger.Debugf("Render(): Cleared isBatchRenderNumColsSet to false (batchRenderNumCols was %d).", t.batchRenderNumCols)
 	}()
 
@@ -1383,9 +1424,10 @@ func (t *Table) render() error {
 		(t.caption.Spot >= tw.SpotTopLeft && t.caption.Spot <= tw.SpotBottomRight)
 
 	var tableStringBuffer *strings.Builder
-	targetWriter := t.writer
-	originalWriter := t.writer // Save original writer for restoration if needed
+	targetWriter := t.writer // Can be the original writer or the MultiWriter.
 
+	// If a caption is present, the main table content must be rendered to an
+	// in-memory buffer first to calculate its final width.
 	if isTopOrBottomCaption {
 		tableStringBuffer = &strings.Builder{}
 		targetWriter = tableStringBuffer
@@ -1394,17 +1436,15 @@ func (t *Table) render() error {
 		t.logger.Debugf("No caption detected. Rendering table core directly to writer.")
 	}
 
-	// Render Table Core
+	// Point the table's writer to the target (either the final destination or the buffer).
 	t.writer = targetWriter
 	ctx, mctx, err := t.prepareContexts()
 	if err != nil {
-		t.writer = originalWriter
 		t.logger.Errorf("prepareContexts failed: %v", err)
 		return errors.Newf("failed to prepare table contexts").Wrap(err)
 	}
 
 	if err := ctx.renderer.Start(t.writer); err != nil {
-		t.writer = originalWriter
 		t.logger.Errorf("Renderer Start() error: %v", err)
 		return errors.Newf("renderer start failed").Wrap(err)
 	}
@@ -1436,18 +1476,21 @@ func (t *Table) render() error {
 		renderError = true
 	}
 
-	t.writer = originalWriter // Restore original writer
+	// Restore the writer to the original for the caption-handling logic.
+	// This is necessary because the caption must be written to the final
+	// destination, not the temporary buffer used for the table body.
+	t.writer = originalWriter
 
 	if renderError {
-		return firstRenderErr // Return error from core rendering if any
+		return firstRenderErr
 	}
 
-	// Caption Handling & Final Output ---
+	// Caption Handling & Final Output
 	if isTopOrBottomCaption {
 		renderedTableContent := tableStringBuffer.String()
 		t.logger.Debugf("[Render] Table core buffer length: %d", len(renderedTableContent))
 
-		// Check if the buffer is empty AND borders are enabled
+		// Handle edge case where table is empty but should have borders.
 		shouldHaveBorders := t.renderer != nil && (t.renderer.Config().Borders.Top.Enabled() || t.renderer.Config().Borders.Bottom.Enabled())
 		if len(renderedTableContent) == 0 && shouldHaveBorders {
 			var sb strings.Builder
@@ -1499,7 +1542,7 @@ func (t *Table) render() error {
 
 	t.hasPrinted = true
 	t.logger.Info("Render() completed.")
-	return nil // Success
+	return nil
 }
 
 // renderFooter renders the table's footer section with borders and padding.

tests/feature_test.go

@@ -2,6 +2,7 @@ package tests
 
 import (
 	"bytes"
+	"io"
 	"testing"
 
 	"github.com/olekukonko/tablewriter"
@@ -253,3 +254,129 @@ func TestTrimLine(t *testing.T) {
 		t.Error(table.Debug())
 	}
 }
+
+// A simple ByteCounter to demonstrate a custom counter implementation.
+type ByteCounter struct {
+	count int
+}
+
+func (bc *ByteCounter) Write(p []byte) (n int, err error) {
+	bc.count += len(p)
+	return len(p), nil
+}
+func (bc *ByteCounter) Total() int {
+	return bc.count
+}
+
+// TestLinesCounter verifies the functionality of the WithLineCounter and WithCounters options.
+func TestLinesCounter(t *testing.T) {
+
+	data := [][]string{
+		{"A", "The Good", "500"},
+		{"B", "The Very Very Bad Man", "288"},
+		{"C", "The Ugly", "120"},
+		{"D", "The Gopher", "800"},
+	}
+
+	// Test Case 1: Default line counting on a standard table using the new API.
+	t.Run("WithLineCounter", func(t *testing.T) {
+		table := tablewriter.NewTable(io.Discard,
+			tablewriter.WithLineCounter(), // Use the new, explicit function.
+		)
+		table.Header("Name", "Sign", "Rating")
+		table.Bulk(data)
+		table.Render()
+
+		// Expected: 1 Top border + 1 Header + 1 Separator + 4 Rows + 1 Bottom border = 8
+		expectedLines := 8
+		if got := table.Lines(); got != expectedLines {
+			t.Errorf("expected %d lines, but got %d", expectedLines, got)
+		}
+	})
+
+	// Test Case 2: Line counting with auto-wrapping enabled.
+	t.Run("LineCounterWithWrapping", func(t *testing.T) {
+		table := tablewriter.NewTable(io.Discard,
+			tablewriter.WithLineCounter(), // Use the new, explicit function.
+			tablewriter.WithRowAutoWrap(tw.WrapNormal),
+			tablewriter.WithMaxWidth(40),
+		)
+		table.Header("Name", "Sign", "Rating")
+		table.Bulk(data)
+		table.Render()
+
+		// Expected: 1 Top border + 1 Header + 1 Separator + 1+3+1+1 Rows + 1 Bottom border = 10
+		expectedLines := 10
+		if got := table.Lines(); got != expectedLines {
+			t.Errorf("expected %d lines with wrapping, but got %d", expectedLines, got)
+		}
+	})
+
+	// Test Case 3: Ensure Lines() returns -1 when no counter is enabled at all.
+	t.Run("NoCounters", func(t *testing.T) {
+		table := tablewriter.NewTable(io.Discard) // No counter options
+		table.Header("Name", "Sign")
+		table.Append("A", "B")
+		table.Render()
+
+		expected := -1
+		if got := table.Lines(); got != expected {
+			t.Errorf("expected %d when no counter is used, but got %d", expected, got)
+		}
+	})
+
+	// Test Case 4: Use a custom counter and verify it's retrieved via Counters().
+	t.Run("WithCustomCounter", func(t *testing.T) {
+		byteCounter := &ByteCounter{}
+		var buf bytes.Buffer
+
+		table := tablewriter.NewTable(&buf,
+			tablewriter.WithCounters(byteCounter), // Use the new plural function for custom counters.
+		)
+		table.Header("A", "B")
+		table.Append("1", "2")
+		table.Render()
+
+		// Crucial Test: Lines() should return -1 because no *LineCounter* was added.
+		if got := table.Lines(); got != -1 {
+			t.Errorf("expected Lines() to return -1 when only a custom counter is used, but got %d", got)
+		}
+
+		// Verify the custom counter via the Counters() method.
+		allCounters := table.Counters()
+		if len(allCounters) != 1 {
+			t.Fatalf("expected 1 counter, but found %d", len(allCounters))
+		}
+
+		if custom, ok := allCounters[0].(*ByteCounter); ok {
+			if custom.Total() <= 0 {
+				t.Errorf("expected a positive byte count from custom counter, but got %d", custom.Total())
+			}
+			if custom.Total() != buf.Len() {
+				t.Errorf("byte counter total (%d) does not match buffer length (%d)", custom.Total(), buf.Len())
+			}
+		} else {
+			t.Error("expected the first counter to be of type *ByteCounter")
+		}
+	})
+
+	// Test Case 5: Ensure Lines() finds the line counter even when mixed with others.
+	t.Run("LinesWithMixedCounters", func(t *testing.T) {
+		byteCounter := &ByteCounter{}
+
+		// Add counters in a specific order: custom first, then default.
+		table := tablewriter.NewTable(io.Discard,
+			tablewriter.WithCounters(byteCounter),
+			tablewriter.WithLineCounter(),
+		)
+		table.Header("Name", "Sign", "Rating")
+		table.Bulk(data)
+		table.Render()
+
+		// Lines() should still find the line count correctly, regardless of order.
+		expectedLines := 8
+		if got := table.Lines(); got != expectedLines {
+			t.Errorf("expected %d lines even with mixed counters, but got %d", expectedLines, got)
+		}
+	})
+}