diskspace/dir_entry.go

package main

import (
	"errors"
	"os"
	"path/filepath"
	"sort"
	"strings"
)

type DirectoryEntry struct {
	path    string
	size    *uint64
	modSize *uint64
	entries []Entry
	parent  Entry

	removal bool
	noPerm  bool
	isMount bool
}

var _ Entry = &DirectoryEntry{}

func (e *DirectoryEntry) Name() string {
	return filepath.Base(e.path)
}

func (e *DirectoryEntry) Path() string {
	return e.path
}

func (e *DirectoryEntry) Size() uint64 {
	if e.size == nil {
		size := uint64(0)
		for _, c := range e.Entries() {
			size += c.Size()
		}
		e.size = &size

		e.modSize = new(uint64)
		*e.modSize = *e.size
	}

	return *e.size
}

func (e *DirectoryEntry) Entries() []Entry {
	if e.entries == nil {
		e.Scan(nil, nil)
	}
	return e.entries
}

func (e *DirectoryEntry) String() string {
	if e.removal {
		return ColorRed.Sprintf("%s %s (%s)", icon("folder"), e.Name(), fmtSize(e.SizeModified()))
	}

	if e.Size() == e.SizeModified() {
		return ColorBlue.Sprintf("%s %s (%s)", icon("folder"), e.Name(), fmtSize(e.Size()))
	} else {
		return ColorBlue.Sprintf("%s %s (%s / %s)", icon("folder"), e.Name(), ColorRed.Sprintf(fmtSize(e.SizeModified())), fmtSize(e.Size()))
	}
}

func (e *DirectoryEntry) stringRecursive(depth, maxDepth int, b *strings.Builder) {
	if depth > maxDepth {
		return
	}

	b.WriteString(strings.Repeat("  ", depth))
	b.WriteString(e.String())
	b.WriteRune('\n')

	for _, entry := range e.Entries() {
		entry.stringRecursive(depth+1, maxDepth, b)
	}
}

func (e *DirectoryEntry) StringRecursive(depth int) string {
	b := new(strings.Builder)
	e.stringRecursive(0, depth, b)
	return b.String()
}

func (e *DirectoryEntry) Scan(ch chan<- string, mounts map[string]struct{}) {
	dirEntries, err := os.ReadDir(e.Path())
	if errors.Is(err, os.ErrPermission) {
		e.noPerm = true
		e.entries = []Entry{}
		return
	} else if err != nil {
		e.noPerm = true
		e.entries = []Entry{}
		return
	}

	e.entries = make([]Entry, 0, len(dirEntries))
	for _, dirEntry := range dirEntries {
		entry := NewEntryFromDirEntry(e, dirEntry)
		if _, isMount := mounts[entry.Path()]; !isMount {
			e.entries = append(e.entries, entry)
			entry.Scan(ch, mounts)
		}
	}

	if ch != nil {
		ch <- e.Path()
	}
	sort.Slice(e.entries, func(i, j int) bool {
		return e.entries[i].SizeModified() < e.entries[j].SizeModified()
	})
}

func (e *DirectoryEntry) Parent() Entry {
	return e.parent
}

func (e *DirectoryEntry) Entry(path string) Entry {
	if path == "." || path == "" {
		return e
	}

	splits := strings.Split(path, "/")
	for _, entry := range e.Entries() {
		if entry.Name() == splits[0] {
			return entry.Entry(strings.Join(splits[1:], "/"))
		}
	}

	return nil
}

func (e *DirectoryEntry) IsDir() bool {
	return true
}

func (e *DirectoryEntry) SetRemovalMark(mark bool) {
	e.removal = mark
	for _, entry := range e.Entries() {
		entry.SetRemovalMark(mark)
	}

	if e.Parent() != nil {
		pentries := e.Parent().Entries()
		sort.Slice(pentries, func(i, j int) bool {
			return pentries[i].SizeModified() < pentries[j].SizeModified()
		})
	}
}

func (e *DirectoryEntry) RemovalMark() bool {
	return e.removal
}

func (e *DirectoryEntry) SizeModified() uint64 {
	if e.size == nil {
		e.Size()
	}

	return *e.modSize
}

func (e *DirectoryEntry) modifySize(modifier uint64) {
	*e.modSize += modifier
	if e.parent != nil {
		e.parent.modifySize(modifier)
	}
}

func (e *DirectoryEntry) RemovalStats(stats *RemovalStats) {
	if e.RemovalMark() {
		stats.Dirs++
		if e.Parent() != nil && !e.Parent().RemovalMark() {
			stats.Paths = append(stats.Paths, e.Path())
		}
	}

	for _, entry := range e.Entries() {
		entry.RemovalStats(stats)
	}
}