Add --show-tokens to command line input

Author: yuchenliu15

cmd/grab/main.go

@@ -47,6 +47,7 @@ func main() {
 	var resolveDeps bool
 	var maxDepth int
 	var maxFileSizeStr string
+	var showTokenCount bool
 
 	flag.BoolVar(&showHelp, "help", false, "Display help information")
 	flag.BoolVar(&showHelp, "h", false, "Display help information (shorthand)")
@@ -85,6 +86,8 @@ func main() {
 	maxFileSizeUsage := "Maximum file size to include (e.g., 50kb, 2MB). No limit by default."
 	flag.StringVar(&maxFileSizeStr, "max-file-size", "", maxFileSizeUsage)
 
+	flag.BoolVar(&showTokenCount, "show-tokens", false, "Show the number of tokens for each file")
+
 	flag.Parse()
 
 	if showHelp {
@@ -154,15 +157,16 @@ func main() {
 		runNonInteractive(root, filterMgr, outputPath, useTempFile, formatName, skipRedaction, resolveDeps, maxDepth, maxFileSize)
 	} else {
 		config := model.Config{
-			RootPath:      root,
-			FilterMgr:     filterMgr,
-			OutputPath:    outputPath,
-			UseTempFile:   useTempFile,
-			Format:        formatName,
-			SkipRedaction: skipRedaction,
-			ResolveDeps:   resolveDeps,
-			MaxDepth:      maxDepth,
-			MaxFileSize:   maxFileSize,
+			RootPath:       root,
+			FilterMgr:      filterMgr,
+			OutputPath:     outputPath,
+			UseTempFile:    useTempFile,
+			Format:         formatName,
+			SkipRedaction:  skipRedaction,
+			ResolveDeps:    resolveDeps,
+			MaxDepth:       maxDepth,
+			MaxFileSize:    maxFileSize,
+			ShowTokenCount: showTokenCount,
 		}
 
 		m := model.NewModel(config)

internal/model/model.go

@@ -56,18 +56,20 @@ type Model struct {
 	isGrabbing        bool
 	redactSecrets     bool
 	resolveDeps       bool
+	showTokenCount    bool
 }
 
 type Config struct {
-	FilterMgr     *filesystem.FilterManager
-	RootPath      string
-	OutputPath    string
-	Format        string
-	MaxDepth      int
-	MaxFileSize   int64
-	UseTempFile   bool
-	SkipRedaction bool
-	ResolveDeps   bool
+	FilterMgr      *filesystem.FilterManager
+	RootPath       string
+	OutputPath     string
+	Format         string
+	MaxDepth       int
+	MaxFileSize    int64
+	UseTempFile    bool
+	SkipRedaction  bool
+	ResolveDeps    bool
+	ShowTokenCount bool
 }
 
 func NewModel(config Config) Model {
@@ -105,6 +107,7 @@ func NewModel(config Config) Model {
 			Width:  80,
 			Height: 10,
 		},
-		cursor: 0,
+		cursor:         0,
+		showTokenCount: config.ShowTokenCount,
 	}
 }

internal/model/view.go

@@ -2,9 +2,9 @@ package model
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"strings"
-	"os"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/epilande/codegrab/internal/ui"
@@ -235,11 +235,13 @@ func (m *Model) refreshViewportContent() {
 			if node.IsDependency {
 				rawSuffix += " [dep]"
 			}
-			if ok, err := utils.IsTextFile(node.Path); ok && err == nil {
-				if contentBytes, err := os.ReadFile(node.Path); err == nil {
-					content := string(contentBytes)
-					tokensEstimate := utils.EstimateTokens(content)
-					rawSuffix += fmt.Sprintf(" [%d tokens]", tokensEstimate)
+			if m.showTokenCount {
+				if ok, err := utils.IsTextFile(node.Path); ok && err == nil {
+					if contentBytes, err := os.ReadFile(node.Path); err == nil {
+						content := string(contentBytes)
+						tokensEstimate := utils.EstimateTokens(content)
+						rawSuffix += fmt.Sprintf(" [%d tokens]", tokensEstimate)
+					}
 				}
 			}
 		}