Merge pull request #13 from yuchenliu15/add_token_counts

Add token counts in view

Author: epilande

README.md

@@ -104,6 +104,7 @@ grab [options] [directory]
 | `--max-depth <depth>`    | Maximum depth for dependency resolution (`-1` for unlimited, default: `1`). Only effective with `--deps`.                                                                                            |
 | `--max-file-size <size>` | Maximum file size to include (e.g., `"100kb"`, `"2MB"`). No limit by default. Files exceeding the specified size will be skipped.                                                                    |
 | `--theme <name>`         | Set the UI theme. Available: catppuccin-latte, catppuccin-frappe, catppuccin-macchiato, catppuccin-mocha, rose-pine, rose-pine-dawn, rose-pine-moon, dracula, nord. (default: `"catppuccin-mocha"`). |
+| `--show-tokens`          | Show the number of tokens for each file in file tree. |
 
 ### 📖 Examples
 

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,11 +2,13 @@ package model
 
 import (
 	"fmt"
+	"os"
 	"path/filepath"
 	"strings"
 
 	"github.com/charmbracelet/lipgloss"
 	"github.com/epilande/codegrab/internal/ui"
+	"github.com/epilande/codegrab/internal/utils"
 )
 
 func (m Model) View() string {
@@ -229,10 +231,20 @@ func (m *Model) refreshViewportContent() {
 			if count := dirSelectedCounts[node.Path]; count > 0 {
 				rawSuffix = fmt.Sprintf(" [%d]", count)
 			}
-		} else if node.IsDependency {
-			rawSuffix = " [dep]"
+		} else {
+			if node.IsDependency {
+				rawSuffix += " [dep]"
+			}
+			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)
+					}
+				}
+			}
 		}
-
 		isCursorLine := i == m.cursor
 
 		rendered := ui.StyleFileLine(

internal/ui/help.go

@@ -54,6 +54,7 @@ const UsageText = `Usage:
     --theme <name>           Set the UI theme. Available: catppuccin-latte, catppuccin-frappe,
                              catppuccin-macchiato, catppuccin-mocha, rose-pine, rose-pine-dawn,
 	                           rose-pine-moon, dracula, nord. (default: "catppuccin-mocha").
+    --show-tokens            Show the number of tokens for each file in file tree.
 
   Examples:
     # Run interactively in the current directory