diff --git a/walker.go b/walker.go index bfec609b..a44950fd 100644 --- a/walker.go +++ b/walker.go @@ -15,6 +15,7 @@ import ( type WalkOpt struct { IncludePaths []string // todo: remove? ExcludePatterns []string + Map func(*Stat) bool } func Walk(ctx context.Context, p string, opt *WalkOpt, fn filepath.WalkFunc) error { @@ -138,7 +139,12 @@ func Walk(ctx context.Context, p string, opt *WalkOpt, fn filepath.WalkFunc) err case <-ctx.Done(): return ctx.Err() default: - if err := fn(path, &StatInfo{stat}, nil); err != nil { + if opt != nil && opt.Map != nil { + if allowed := opt.Map(stat); !allowed { + return nil + } + } + if err := fn(stat.Path, &StatInfo{stat}, nil); err != nil { return err } } diff --git a/walker_test.go b/walker_test.go index a432afb1..9572316f 100644 --- a/walker_test.go +++ b/walker_test.go @@ -6,6 +6,7 @@ import ( "io/ioutil" "os" "path/filepath" + "strings" "testing" "github.com/pkg/errors" @@ -72,6 +73,33 @@ file foo/bar2 } +func TestWalkerMap(t *testing.T) { + d, err := tmpDir(changeStream([]string{ + "ADD bar file", + "ADD foo dir", + "ADD foo2 file", + "ADD foo/bar2 file", + })) + assert.NoError(t, err) + defer os.RemoveAll(d) + b := &bytes.Buffer{} + err = Walk(context.Background(), d, &WalkOpt{ + Map: func(s *Stat) bool { + if strings.HasPrefix(s.Path, "foo") { + s.Path = "_" + s.Path + return true + } + return false + }, + }, bufWalk(b)) + assert.NoError(t, err) + + assert.Equal(t, `dir _foo +file _foo/bar2 +file _foo2 +`, string(b.Bytes())) +} + func bufWalk(buf *bytes.Buffer) filepath.WalkFunc { return func(path string, fi os.FileInfo, err error) error { stat, ok := fi.Sys().(*Stat)