refactor walker

This commit is contained in:
Dominic Breuker
2018-03-12 08:40:30 +01:00
parent f2783e95c6
commit cc846bdad1

View File

@@ -7,53 +7,78 @@ import (
"path/filepath" "path/filepath"
) )
type Walker struct{} type Walker struct {
}
func NewWalker() *Walker { func NewWalker() *Walker {
return &Walker{} return &Walker{}
} }
type chans struct {
dirCh chan string
errCh chan error
doneCh chan struct{}
}
func newChans() *chans {
return &chans{
dirCh: make(chan string),
errCh: make(chan error),
doneCh: make(chan struct{}),
}
}
const maxInt = int(^uint(0) >> 1) const maxInt = int(^uint(0) >> 1)
func (w *Walker) Walk(root string, depth int) (dirCh chan string, errCh chan error, doneCh chan struct{}) { func (w *Walker) Walk(root string, depth int) (dirCh chan string, errCh chan error, doneCh chan struct{}) {
if depth < 0 { if depth < 0 {
depth = maxInt depth = maxInt
} }
dirCh = make(chan string) c := newChans()
errCh = make(chan error)
doneCh = make(chan struct{})
go func() { go func() {
defer close(dirCh) defer close(dirCh)
descent(root, depth-1, dirCh, errCh, doneCh) descent(root, depth-1, c)
}() }()
return dirCh, errCh, doneCh return c.dirCh, c.errCh, c.doneCh
} }
func descent(dir string, depth int, dirCh chan string, errCh chan error, doneCh chan struct{}) { func descent(dir string, depth int, c *chans) {
_, err := os.Stat(dir) if done := emitDir(dir, depth, c); done {
if err != nil {
errCh <- fmt.Errorf("visiting %s: %v", dir, err)
return
}
select {
case dirCh <- dir:
case <-doneCh:
return
}
if depth < 0 {
return return
} }
handleSubDirs(dir, depth, c)
}
func emitDir(dir string, depth int, c *chans) bool {
_, err := os.Stat(dir)
if err != nil {
c.errCh <- fmt.Errorf("visiting %s: %v", dir, err)
return true
}
select {
case c.dirCh <- dir:
case <-c.doneCh:
return true
}
if depth < 0 {
return true
}
return false
}
func handleSubDirs(dir string, depth int, c *chans) {
ls, err := ioutil.ReadDir(dir) ls, err := ioutil.ReadDir(dir)
if err != nil { if err != nil {
errCh <- fmt.Errorf("opening dir %s: %v", dir, err) c.errCh <- fmt.Errorf("opening dir %s: %v", dir, err)
} }
for _, e := range ls { for _, e := range ls {
if e.IsDir() { if e.IsDir() {
newDir := filepath.Join(dir, e.Name()) newDir := filepath.Join(dir, e.Name())
descent(newDir, depth-1, dirCh, errCh, doneCh) descent(newDir, depth-1, c)
} }
} }
} }