Merge pull request #7 from noborus/add-align

Add align option

Author: noborus

cmd/root.go

@@ -3,6 +3,7 @@ package cmd
 import (
 	"fmt"
 	"os"
+	"strings"
 
 	"github.com/noborus/guesswidth"
 	"github.com/spf13/cobra"
@@ -16,7 +17,7 @@ var rootCmd = &cobra.Command{
 split them, insert fences and output.`,
 	Version: guesswidth.Version(),
 	Run: func(cmd *cobra.Command, args []string) {
-		writeTable(args)
+		writeTable()
 	},
 }
 
@@ -25,16 +26,21 @@ var (
 	header     int
 	limitSplit int
 	scanNum    int
+	align      bool
 )
 
-func writeTable(args []string) {
+func writeTable() {
 	g := guesswidth.NewReader(os.Stdin)
 	g.Header = header - 1
 	g.LimitSplit = limitSplit
 	g.TrimSpace = false
 	if scanNum > 0 {
 		g.ScanNum = scanNum
 	}
+	if align {
+		writeAlign(g)
+		return
+	}
 	write(g)
 }
 
@@ -54,6 +60,27 @@ func write(g *guesswidth.GuessWidth) {
 	}
 }
 
+func writeAlign(g *guesswidth.GuessWidth) {
+	for _, row := range g.ReadAll() {
+		for n, col := range row {
+			if n > 0 {
+				fmt.Print(fence)
+			}
+			col = strings.TrimSpace(col)
+			if g.Widths[n].Justified == guesswidth.Right {
+				fmt.Printf("%*s", g.Widths[n].Width, col)
+			} else {
+				if len(g.Widths)-1 == n {
+					fmt.Printf("%s", col)
+				} else {
+					fmt.Printf("%-*s", g.Widths[n].Width, col)
+				}
+			}
+		}
+		fmt.Println()
+	}
+}
+
 // Execute adds all child commands to the root command and sets flags appropriately.
 // This is called by main.main(). It only needs to happen once to the rootCmd.
 func Execute() {
@@ -66,7 +93,7 @@ func init() {
 	rootCmd.PersistentFlags().IntVar(&header, "header", 1, "header line number")
 	rootCmd.PersistentFlags().IntVar(&limitSplit, "split", -1, "maximum number of splits")
 	rootCmd.PersistentFlags().IntVar(&scanNum, "scannum", 100, "number of line to scan")
-	rootCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
+	rootCmd.PersistentFlags().BoolVarP(&align, "align", "a", false, "align the output")
 }
 
 // initConfig reads in config file and ENV variables if set.

guesswidth.go

@@ -21,6 +21,8 @@ type GuessWidth struct {
 	reader *bufio.Reader
 	// pos is a list of separator positions.
 	pos []int
+	// Widths is the width of the column.
+	Widths []Cols
 	// preLines stores the lines read for scan.
 	preLines []string
 	// preCount is the number returned by read.
@@ -39,6 +41,17 @@ type GuessWidth struct {
 	TrimSpace bool
 }
 
+type Cols struct {
+	Width      int
+	Justified  int
+	rightCount int
+}
+
+const (
+	Left = iota
+	Right
+)
+
 // NewReader returns a new Reader that reads from r.
 func NewReader(r io.Reader) *GuessWidth {
 	reader := bufio.NewReader(r)
@@ -61,17 +74,69 @@ func (g *GuessWidth) ReadAll() [][]string {
 		g.Scan(g.ScanNum)
 	}
 
+	g.Widths = make([]Cols, len(g.pos)+1)
 	var rows [][]string
 	for {
 		columns, err := g.Read()
 		if err != nil {
 			break
 		}
+		g.UpdateMaxWidth(columns)
 		rows = append(rows, columns)
 	}
+
+	g.SetJustified(len(rows) / 2)
 	return rows
 }
 
+// UpdateMaxWidth updates the maximum width of the column.
+func (g *GuessWidth) UpdateMaxWidth(columns []string) []Cols {
+	if len(g.Widths) < len(columns) {
+		for n := len(g.Widths); n < len(columns); n++ {
+			g.Widths = append(g.Widths, Cols{})
+		}
+	}
+
+	for n, col := range columns {
+		width := runewidth.StringWidth(col)
+		if width > g.Widths[n].Width {
+			g.Widths[n].Width = width
+		}
+		if isRightAlign(col) {
+			g.Widths[n].rightCount++
+		}
+	}
+	return g.Widths
+}
+
+// SetJustified sets the justification of the column.
+func (g *GuessWidth) SetJustified(threshold int) []Cols {
+	for n, col := range g.Widths {
+		if col.rightCount < threshold {
+			col.Justified = Left
+		} else {
+			col.Justified = Right
+		}
+		g.Widths[n] = col
+	}
+	return g.Widths
+}
+
+func isRightAlign(str string) bool {
+	if str == "" {
+		return false
+	}
+	for n := 0; n < len(str); n++ {
+		if str[n] != ' ' {
+			return false
+		}
+		if str[len(str)-n-1] != ' ' {
+			return true
+		}
+	}
+	return false
+}
+
 // Scan preReads and parses the lines.
 func (g *GuessWidth) Scan(num int) {
 	for i := 0; i < num; i++ {

guesswidth_test.go

@@ -23,6 +23,7 @@ func TestGuessWidth_ReadAll(t *testing.T) {
 		name   string
 		fields fields
 		want   [][]string
+		want2  []Cols
 	}{
 		{
 			name: "ps",
@@ -38,6 +39,12 @@ func TestGuessWidth_ReadAll(t *testing.T) {
 				{"302965", " pts/3   ", " 00:00:11", "zsh"},
 				{"709737", " pts/3   ", " 00:00:00", "ps"},
 			},
+			want2: []Cols{
+				{6, 1, 1},
+				{9, 0, 0},
+				{9, 1, 3},
+				{3, 0, 0},
+			},
 		},
 		{
 			name: "ps overflow",
@@ -55,6 +62,19 @@ noborus   721971  0.0  0.0  13716  3524 pts/3    R+   10:39   0:00 ps aux`)),
 				{"noborus  ", " 703052", "  2.1", "  0.7", " 1184814400", " 230920", " ?  ", " Sl  ", " 10:03  ", " 0:45", "/opt/google/chrome/chrome"},
 				{"noborus  ", " 721971", "  0.0", "  0.0", "  13716", "  3524", " pts/3   ", " R+  ", " 10:39  ", " 0:00", "ps aux"},
 			},
+			want2: []Cols{
+				{9, 0, 0},
+				{7, 1, 4},
+				{5, 1, 4},
+				{5, 1, 4},
+				{11, 1, 4},
+				{7, 1, 4},
+				{9, 0, 0},
+				{5, 0, 1},
+				{8, 0, 0},
+				{5, 1, 4},
+				{25, 0, 0},
+			},
 		},
 		{
 			name: "ps limit",
@@ -71,6 +91,11 @@ noborus   721971  0.0  0.0  13716  3524 pts/3    R+   10:39   0:00 ps aux`)),
 				{"302965", " pts/3   ", "00:00:11 zsh"},
 				{"709737", " pts/3   ", "00:00:00 ps"},
 			},
+			want2: []Cols{
+				{6, 1, 1},
+				{9, 0, 0},
+				{12, 1, 1},
+			},
 		},
 	}
 	for _, tt := range tests {
@@ -89,6 +114,9 @@ noborus   721971  0.0  0.0  13716  3524 pts/3    R+   10:39   0:00 ps aux`)),
 			if got := g.ReadAll(); !reflect.DeepEqual(got, tt.want) {
 				t.Errorf("GuessWidth.ReadAll() = \n%#v, want \n%#v", got, tt.want)
 			}
+			if got2 := g.Widths; !reflect.DeepEqual(got2, tt.want2) {
+				t.Errorf("GuessWidth.ReadAll() = \n%v, want \n%v", got2, tt.want2)
+			}
 		})
 	}
 }