Added --maximum-files and --minimum-files flags, changed structure for stats struct

README.md

@@ -89,6 +89,8 @@ When accessed, these endpoints return the contents of the index, in HTML and JSO
 
 ## Usage output
 ```
+Serves random images from the specified directories.
+
 Usage:
   roulette <path> [path]... [flags]
 
@@ -99,6 +101,8 @@ Flags:
   -d, --debug                  expose debug endpoint
   -f, --filter                 enable filtering
   -h, --help                   help for roulette
+      --maximum-files uint32   skip directories with file counts over this value (default 4294967295)
+      --minimum-files uint32   skip directories with file counts under this value
   -p, --port uint16            port to listen on (default 8080)
   -r, --recursive              recurse into subdirectories
   -s, --sort                   enable sorting

cmd/files.go

@@ -66,36 +66,49 @@ type ScanStats struct {
 	filesMatched       uint64
 	filesSkipped       uint64
 	directoriesMatched uint64
+	directoriesSkipped uint64
 }
 
 func (s *ScanStats) FilesTotal() uint64 {
 	return atomic.LoadUint64(&s.filesMatched) + atomic.LoadUint64(&s.filesSkipped)
 }
 
-func (s *ScanStats) incrementFilesMatched() {
+func (s *ScanStats) incrementFilesMatched(n int) {
-	atomic.AddUint64(&s.filesMatched, 1)
+	atomic.AddUint64(&s.filesMatched, uint64(n))
 }
 
 func (s *ScanStats) FilesMatched() uint64 {
 	return atomic.LoadUint64(&s.filesMatched)
 }
 
-func (s *ScanStats) incrementFilesSkipped() {
+func (s *ScanStats) incrementFilesSkipped(n int) {
-	atomic.AddUint64(&s.filesSkipped, 1)
+	atomic.AddUint64(&s.filesSkipped, uint64(n))
 }
 
 func (s *ScanStats) FilesSkipped() uint64 {
 	return atomic.LoadUint64(&s.filesSkipped)
 }
 
-func (s *ScanStats) incrementDirectoriesMatched() {
+func (s *ScanStats) DirectoriesTotal() uint64 {
-	atomic.AddUint64(&s.directoriesMatched, 1)
+	return atomic.LoadUint64(&s.directoriesMatched) + atomic.LoadUint64(&s.directoriesSkipped)
+}
+
+func (s *ScanStats) incrementDirectoriesMatched(n int) {
+	atomic.AddUint64(&s.directoriesMatched, uint64(n))
 }
 
 func (s *ScanStats) DirectoriesMatched() uint64 {
 	return atomic.LoadUint64(&s.directoriesMatched)
 }
 
+func (s *ScanStats) incrementDirectoriesSkipped(n int) {
+	atomic.AddUint64(&s.directoriesSkipped, uint64(n))
+}
+
+func (s *ScanStats) DirectoriesSkipped() uint64 {
+	return atomic.LoadUint64(&s.directoriesSkipped)
+}
+
 type Path struct {
 	base      string
 	number    int
@@ -182,7 +195,7 @@ func appendPath(directory, path string, files *Files, stats *ScanStats, shouldCa
 
 	files.Append(directory, path)
 
-	stats.incrementFilesMatched()
+	stats.incrementFilesMatched(1)
 
 	return nil
 }
@@ -205,7 +218,7 @@ func appendPaths(path string, files *Files, filters *Filters, stats *ScanStats)
 				filename,
 				filters.excludes[i],
 			) {
-				stats.incrementFilesSkipped()
+				stats.incrementFilesSkipped(1)
 
 				return nil
 			}
@@ -227,7 +240,7 @@ func appendPaths(path string, files *Files, filters *Filters, stats *ScanStats)
 			}
 		}
 
-		stats.incrementFilesSkipped()
+		stats.incrementFilesSkipped(1)
 
 		return nil
 	}
@@ -442,6 +455,26 @@ func pathHasSupportedFiles(path string) (bool, error) {
 	}
 }
 
+func pathCount(path string) (int, int, error) {
+	directories := 0
+	files := 0
+
+	nodes, err := os.ReadDir(path)
+	if err != nil {
+		return 0, 0, err
+	}
+
+	for _, node := range nodes {
+		if node.IsDir() {
+			directories++
+		} else {
+			files++
+		}
+	}
+
+	return files, directories, nil
+}
+
 func scanPath(path string, files *Files, filters *Filters, stats *ScanStats, concurrency *Concurrency) error {
 	var wg sync.WaitGroup
 
@@ -475,7 +508,20 @@ func scanPath(path string, files *Files, filters *Filters, stats *ScanStats, con
 				}
 			}()
 		case info.IsDir():
-			stats.incrementDirectoriesMatched()
+			files, directories, err := pathCount(p)
+			if err != nil {
+				fmt.Println(err)
+			}
+
+			if files > 0 && (files < int(minimumFileCount) || files > int(maximumFileCount)) {
+				// This count will not otherwise include the parent directory itself, so increment by one
+				stats.incrementDirectoriesSkipped(directories + 1)
+				stats.incrementFilesSkipped(files)
+
+				return filepath.SkipDir
+			}
+
+			stats.incrementDirectoriesMatched(1)
 		}
 
 		return err
@@ -506,6 +552,7 @@ func fileList(paths []string, filters *Filters, sort string, index *Index) ([]st
 		filesMatched:       0,
 		filesSkipped:       0,
 		directoriesMatched: 0,
+		directoriesSkipped: 0,
 	}
 
 	concurrency := &Concurrency{
@@ -540,11 +587,12 @@ func fileList(paths []string, filters *Filters, sort string, index *Index) ([]st
 	fileList = prepareDirectories(files, sort)
 
 	if verbose {
-		fmt.Printf("%s | Indexed %d/%d files across %d directories in %s\n",
+		fmt.Printf("%s | Indexed %d/%d files across %d/%d directories in %s\n",
 			time.Now().Format(LogDate),
 			stats.FilesMatched(),
 			stats.FilesTotal(),
 			stats.DirectoriesMatched(),
+			stats.DirectoriesTotal(),
 			time.Since(startTime),
 		)
 	}

cmd/root.go

@@ -6,12 +6,13 @@ package cmd
 
 import (
 	"log"
+	"math"
 
 	"github.com/spf13/cobra"
 )
 
 const (
-	Version string = "0.53.2"
+	Version string = "0.54.0"
 )
 
 var (
@@ -20,6 +21,8 @@ var (
 	cacheFile        string
 	debug            bool
 	filtering        bool
+	maximumFileCount uint32
+	minimumFileCount uint32
 	port             uint16
 	recursive        bool
 	sorting          bool
@@ -61,6 +64,8 @@ func init() {
 	rootCmd.Flags().StringVar(&cacheFile, "cache-file", "", "path to optional persistent cache file")
 	rootCmd.Flags().BoolVarP(&debug, "debug", "d", false, "expose debug endpoint")
 	rootCmd.Flags().BoolVarP(&filtering, "filter", "f", false, "enable filtering")
+	rootCmd.Flags().Uint32Var(&maximumFileCount, "maximum-files", math.MaxUint32, "skip directories with file counts over this value")
+	rootCmd.Flags().Uint32Var(&minimumFileCount, "minimum-files", 0, "skip directories with file counts under this value")
 	rootCmd.Flags().Uint16VarP(&port, "port", "p", 8080, "port to listen on")
 	rootCmd.Flags().BoolVarP(&recursive, "recursive", "r", false, "recurse into subdirectories")
 	rootCmd.Flags().BoolVarP(&sorting, "sort", "s", false, "enable sorting")