mirror of
https://github.com/DominicBreuker/pspy.git
synced 2025-12-21 11:44:51 +00:00
refactor walker
This commit is contained in:
@@ -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)
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user