wip: add git stage reader

Signed-off-by: Brian McGee <[email protected]>

nix/packages/treefmt/default.nix

@@ -46,7 +46,7 @@ in
       "-X github.com/numtide/treefmt/build.Version=v${version}"
     ];
 
-    nativeBuildInputs =
+    nativeBuildInputs = [ pkgs.git ] ++
       # we need some formatters available for the tests
       import ./formatters.nix pkgs;
 

walk/git.go

@@ -4,14 +4,15 @@ import (
 	"bufio"
 	"context"
 	"fmt"
-	"github.com/charmbracelet/log"
-	"github.com/numtide/treefmt/stats"
-	"golang.org/x/sync/errgroup"
 	"io"
 	"os"
 	"os/exec"
 	"path/filepath"
 	"strings"
+
+	"github.com/charmbracelet/log"
+	"github.com/numtide/treefmt/stats"
+	"golang.org/x/sync/errgroup"
 )
 
 type GitReader struct {
@@ -99,12 +100,12 @@ func (g *GitReader) Close() error {
 	return g.eg.Wait()
 }
 
-func NewGitWorktreeReader(
+func newGitReader(
 	root string,
 	path string,
+	args []string,
 	statz *stats.Stats,
 ) (*GitReader, error) {
-
 	// check if the root is a git repository
 	cmd := exec.Command("git", "rev-parse", "--is-inside-work-tree")
 	cmd.Dir = root
@@ -118,9 +119,25 @@ func NewGitWorktreeReader(
 	return &GitReader{
 		root:  root,
 		path:  path,
-		args:  []string{"ls-files"},
+		args:  args,
 		stats: statz,
 		eg:    &errgroup.Group{},
 		log:   log.WithPrefix("walk[git]"),
 	}, nil
 }
+
+func NewGitStageReader(
+	root string,
+	path string,
+	statz *stats.Stats,
+) (*GitReader, error) {
+	return newGitReader(root, path, []string{"diff", "--name-only", "--cached"}, statz)
+}
+
+func NewGitWorktreeReader(
+	root string,
+	path string,
+	statz *stats.Stats,
+) (*GitReader, error) {
+	return newGitReader(root, path, []string{"ls-files"}, statz)
+}

walk/git_test.go

@@ -33,14 +33,24 @@ func TestGitWorktreeReader(t *testing.T) {
 	)
 	as.NoError(err, "failed to init git repository")
 
-	// get worktree and add everything to it
+	// read empty worktree
+	statz := stats.New()
+	reader, err := walk.NewGitWorktreeReader(tempDir, "", &statz)
+	as.NoError(err)
+
+	files := make([]*walk.File, 8)
+	ctx, cancel := context.WithTimeout(context.Background(), 100*time.Millisecond)
+	n, err := reader.Read(ctx, files)
+	cancel()
+	as.Equal(0, n)
+	as.ErrorIs(err, io.EOF)
+
+	// add everything to the worktree
 	wt, err := repo.Worktree()
 	as.NoError(err, "failed to get git worktree")
 	as.NoError(wt.AddGlob("."))
 
-	statz := stats.New()
-
-	reader, err := walk.NewGitWorktreeReader(tempDir, "", &statz)
+	reader, err = walk.NewGitWorktreeReader(tempDir, "", &statz)
 	as.NoError(err)
 
 	count := 0

walk/walk.go

@@ -20,9 +20,11 @@ type Type int
 
 const (
 	Auto Type = iota
-	Git
-	Filesystem
 	Stdin
+	Filesystem
+	Git
+	GitWorktree
+	GitStage
 
 	BatchSize = 1024
 )
@@ -159,12 +161,15 @@ func NewReader(
 			reader, err = NewReader(Filesystem, root, path, db, statz)
 		}
 		return reader, err
-	case Git:
-		reader, err = NewGitWorktreeReader(root, path, statz)
-	case Filesystem:
-		reader = NewFilesystemReader(root, path, statz, BatchSize)
 	case Stdin:
 		return nil, fmt.Errorf("stdin walk type is not supported")
+	case Filesystem:
+		reader = NewFilesystemReader(root, path, statz, BatchSize)
+	case Git, GitWorktree:
+		reader, err = NewGitWorktreeReader(root, path, statz)
+	case GitStage:
+		reader, err = NewGitStageReader(root, path, statz)
+
 	default:
 		return nil, fmt.Errorf("unknown walk type: %v", walkType)
 	}