From ab1d3624bd68c3ae352bf0871877497fde01efc8 Mon Sep 17 00:00:00 2001
From: Brian McGee <brian@bmcgee.ie>
Date: Wed, 10 Jan 2024 16:07:05 +0000
Subject: [PATCH] feat: support reading paths from stdin

Signed-off-by: Brian McGee <brian@bmcgee.ie>
---
 internal/cli/cli.go         |  1 +
 internal/cli/format.go      |  4 ++-
 internal/cli/format_test.go | 56 +++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 1 deletion(-)

diff --git a/internal/cli/cli.go b/internal/cli/cli.go
index 8cb444f..b8872c9 100644
--- a/internal/cli/cli.go
+++ b/internal/cli/cli.go
@@ -20,6 +20,7 @@ type Format struct {
 	Verbosity             int                `name:"verbose" short:"v" type:"counter" default:"0" env:"LOG_LEVEL" help:"Set the verbosity of logs e.g. -vv."`
 
 	Paths []string `name:"paths" arg:"" type:"path" optional:"" help:"Paths to format. Defaults to formatting the whole tree."`
+	Stdin bool     `help:"Format the context passed in via stdin"`
 }
 
 func (f *Format) Configure() {
diff --git a/internal/cli/format.go b/internal/cli/format.go
index 68694de..33dea89 100644
--- a/internal/cli/format.go
+++ b/internal/cli/format.go
@@ -192,7 +192,9 @@ func (f *Format) Run() error {
 	eg.Go(func() (err error) {
 		var walker walk.Walker
 
-		if len(Cli.Paths) > 0 {
+		if Cli.Stdin {
+			walker, err = walk.NewPathReader(os.Stdin)
+		} else if len(Cli.Paths) > 0 {
 			walker, err = walk.NewPathList(Cli.Paths)
 		} else {
 			walker, err = walk.New(Cli.Walk, Cli.TreeRoot)
diff --git a/internal/cli/format_test.go b/internal/cli/format_test.go
index 2dffcc3..d3a96ed 100644
--- a/internal/cli/format_test.go
+++ b/internal/cli/format_test.go
@@ -448,3 +448,59 @@ func TestPathsArg(t *testing.T) {
 	out, err = cmd(t, "-C", tempDir, "-c", "elm/elm.json", "haskell/Nested/Bar.hs")
 	as.ErrorContains(err, "no such file or directory")
 }
+
+func TestStdIn(t *testing.T) {
+	as := require.New(t)
+
+	// capture current cwd, so we can replace it after the test is finished
+	cwd, err := os.Getwd()
+	as.NoError(err)
+
+	t.Cleanup(func() {
+		// return to the previous working directory
+		as.NoError(os.Chdir(cwd))
+	})
+
+	tempDir := test.TempExamples(t)
+	configPath := filepath.Join(tempDir, "/treefmt.toml")
+
+	// change working directory to temp root
+	as.NoError(os.Chdir(tempDir))
+
+	// basic config
+	config := format.Config{
+		Formatters: map[string]*format.Formatter{
+			"echo": {
+				Command:  "echo",
+				Includes: []string{"*"},
+			},
+		},
+	}
+	test.WriteConfig(t, configPath, config)
+
+	// swap out stdin
+	prevStdIn := os.Stdin
+	stdin, err := os.CreateTemp("", "stdin")
+	as.NoError(err)
+
+	os.Stdin = stdin
+
+	t.Cleanup(func() {
+		os.Stdin = prevStdIn
+		_ = os.Remove(stdin.Name())
+	})
+
+	go func() {
+		_, err := stdin.WriteString(`treefmt.toml
+elm/elm.json
+go/main.go
+`)
+		as.NoError(err, "failed to write to stdin")
+		as.NoError(stdin.Sync())
+		_, _ = stdin.Seek(0, 0)
+	}()
+
+	out, err := cmd(t, "-C", tempDir, "--stdin")
+	as.NoError(err)
+	as.Contains(string(out), fmt.Sprintf("%d files changed", 3))
+}