feat: add a max-batch-size option for formatters

This is part of #334

Author: jfly

cmd/format/format.go

@@ -22,10 +22,6 @@ import (
 	bolt "go.etcd.io/bbolt"
 )
 
-const (
-	BatchSize = 1024
-)
-
 var ErrFailOnChange = errors.New("unexpected changes detected, --fail-on-change is enabled")
 
 func Run(v *viper.Viper, statz *stats.Stats, cmd *cobra.Command, paths []string) error {
@@ -123,7 +119,7 @@ func Run(v *viper.Viper, statz *stats.Stats, cmd *cobra.Command, paths []string)
 	}
 
 	// create a composite formatter which will handle applying the correct formatters to each file we traverse
-	formatter, err := format.NewCompositeFormatter(cfg, statz, BatchSize)
+	formatter, err := format.NewCompositeFormatter(cfg, statz, format.BatchSize)
 	if err != nil {
 		return fmt.Errorf("failed to create composite formatter: %w", err)
 	}
@@ -135,7 +131,7 @@ func Run(v *viper.Viper, statz *stats.Stats, cmd *cobra.Command, paths []string)
 	}
 
 	// start traversing
-	files := make([]*walk.File, BatchSize)
+	files := make([]*walk.File, format.BatchSize)
 
 	var (
 		n                  int

cmd/root_test.go

@@ -2501,6 +2501,35 @@ func TestConcurrentInvocation(t *testing.T) {
 	as.NoError(eg.Wait())
 }
 
+func TestMaxBatchSize(t *testing.T) {
+	tempDir := test.TempExamples(t)
+	configPath := filepath.Join(tempDir, "/treefmt.toml")
+
+	test.ChangeWorkDir(t, tempDir)
+
+	maxBatchSize := 1
+	cfg := &config.Config{
+		FormatterConfigs: map[string]*config.Formatter{
+			"echo": {
+				Command:      "test-fmt-only-one-file-at-a-time",
+				Includes:     []string{"*"},
+				MaxBatchSize: &maxBatchSize,
+			},
+		},
+	}
+
+	treefmt(t,
+		withConfig(configPath, cfg),
+		withNoError(t),
+		withStats(t, map[stats.Type]int{
+			stats.Traversed: 33,
+			stats.Matched:   33,
+			stats.Formatted: 33,
+			stats.Changed:   33,
+		}),
+	)
+}
+
 type options struct {
 	args []string
 	env  map[string]string

config/config.go

@@ -62,6 +62,8 @@ type Formatter struct {
 	Excludes []string `mapstructure:"excludes,omitempty" toml:"excludes,omitempty"`
 	// Indicates the order of precedence when executing this Formatter in a sequence of Formatters.
 	Priority int `mapstructure:"priority,omitempty" toml:"priority,omitempty"`
+	// The maximum number of files we should pass to this Formatter at once.
+	MaxBatchSize *int `mapstructure:"max-batch-size" toml:"max-batch-size"`
 }
 
 // SetFlags appends our flags to the provided flag set.

format/formatter.go

@@ -20,6 +20,10 @@ import (
 	"mvdan.cc/sh/v3/interp"
 )
 
+const (
+	BatchSize = 1024
+)
+
 var (
 	ErrInvalidName = errors.New("formatter name must only contain alphanumeric characters, `_` or `-`")
 	// ErrCommandNotFound is returned when the Command for a Formatter is not available.
@@ -46,6 +50,14 @@ func (f *Formatter) Name() string {
 	return f.name
 }
 
+func (f *Formatter) MaxBatchSize() int {
+	if f.config.MaxBatchSize == nil {
+		return BatchSize
+	}
+
+	return *f.config.MaxBatchSize
+}
+
 func (f *Formatter) Priority() int {
 	return f.config.Priority
 }
@@ -78,6 +90,10 @@ func (f *Formatter) Hash(h hash.Hash) error {
 }
 
 func (f *Formatter) Apply(ctx context.Context, files []*walk.File) error {
+	if len(files) > f.MaxBatchSize() {
+		return fmt.Errorf("formatter cannot format %d files at once (max batch size: %d)", len(files), f.MaxBatchSize())
+	}
+
 	start := time.Now()
 
 	// construct args, starting with config

format/scheduler.go

@@ -145,8 +145,10 @@ func (s *scheduler) schedule(ctx context.Context, key batchKey, batch []*walk.Fi
 		for _, name := range key.sequence() {
 			formatter := s.formatters[name]
 
-			if err := formatter.Apply(ctx, batch); err != nil {
-				formatErrors = append(formatErrors, err)
+			for chunk := range slices.Chunk(batch, formatter.MaxBatchSize()) {
+				if err := formatter.Apply(ctx, chunk); err != nil {
+					formatErrors = append(formatErrors, err)
+				}
 			}
 		}
 

nix/packages/treefmt/formatters.nix

@@ -58,4 +58,15 @@ with pkgs; [
       test-fmt-append "$@"
     '';
   })
+  (pkgs.writeShellApplication {
+    name = "test-fmt-only-one-file-at-a-time";
+    text = ''
+      if [ $# -ne 1 ]; then
+        echo "I only support formatting exactly 1 file at a time"
+        exit 1
+      fi
+
+      test-fmt-append "suffix" "$1"
+    '';
+  })
 ]