From 9867d57cf06332ce5259e6ea9c94f45492df46f5 Mon Sep 17 00:00:00 2001
From: myaaaaaaaaa <103326468+myaaaaaaaaa@users.noreply.github.com>
Date: Sun, 24 Nov 2024 04:05:52 -0500
Subject: [PATCH] implement $find

---
 jqx/main.go  |  1 +
 prog.go      | 21 +++++++++++++++++++++
 prog_test.go | 25 ++++++++++++++++++++++++-
 3 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/jqx/main.go b/jqx/main.go
index 7b8641c..891b159 100644
--- a/jqx/main.go
+++ b/jqx/main.go
@@ -35,6 +35,7 @@ func main() {
 		Args: os.Args[1:],
 
 		Open: func(f string) (fs.File, error) { return os.Open(f) },
+		Find: os.DirFS,
 
 		Stdin:  os.Stdin,
 		Stdout: os.Stdout,
diff --git a/prog.go b/prog.go
index 10d1d56..10f439c 100644
--- a/prog.go
+++ b/prog.go
@@ -49,6 +49,7 @@ type Program struct {
 	Args []string
 
 	Open  func(string) (fs.File, error)
+	Find  func(string) fs.FS
 	OutFS fs.FS
 
 	Stdin  io.Reader
@@ -66,6 +67,8 @@ type flags struct {
 	rawIn   bool
 	jsonOut bool
 	env     bool
+
+	find string
 }
 
 func (f *flags) populate(args []string) {
@@ -76,6 +79,8 @@ func (f *flags) populate(args []string) {
 	fset.BoolVar(&f.jsonOut, "j", false, `(json) always output json (strings are unwrapped by default)`)
 	fset.BoolVar(&f.env, "e", false, `(env) enable $env`)
 
+	fset.StringVar(&f.find, "find", "", `enable $find`)
+
 	usage := fset.Usage
 	fset.Usage = func() {
 		usage()
@@ -141,6 +146,22 @@ func (p *Program) Main() (rtErr error) {
 		}
 		state.Globals["$env"] = envVars
 	}
+	if f.find != "" {
+		find := map[string]any{}
+		err := fs.WalkDir(p.Find(f.find), ".", func(path string, d fs.DirEntry, err error) error {
+			if path == "." {
+				return err
+			}
+
+			if d.IsDir() {
+				path = path + "/"
+			}
+			find[path] = true
+			return err
+		})
+		failif(err, "finding subdirs")
+		state.Globals["$find"] = find
+	}
 	query := state.Compile(constString(f.script))
 	for v := range input {
 		for v := range query(v) {
diff --git a/prog_test.go b/prog_test.go
index 6e28a0c..8d63c15 100644
--- a/prog_test.go
+++ b/prog_test.go
@@ -3,8 +3,10 @@ package jqx
 import (
 	"bytes"
 	"io/fs"
+	"slices"
 	"strings"
 	"testing"
+	"testing/fstest"
 )
 
 func testRun(t *testing.T, stdin, want string, p *Program) {
@@ -56,7 +58,7 @@ func TestProgram(t *testing.T) {
 	testRun(t, `"XYZ"`, "_____", &Program{Args: []string{"-e", "$env[.]"}})
 	testRun(t, `"XYZ"`, "error", &Program{Args: []string{"$env[.]"}})
 }
-func TestFS(t *testing.T) {
+func TestOpen(t *testing.T) {
 	testFiles := map[string]any{
 		"a.json":    "[1][2][3]",
 		"b.json":    "[1,2,3]",
@@ -86,6 +88,27 @@ func TestFS(t *testing.T) {
 	p.Args = []string{"-r", ".[][] | .+.", "d.txt", "e.txt"}
 	testRun(t, "", "foofoo qq ww ee rr tt yy", &p)
 }
+
+func TestFS(t *testing.T) {
+	testfs := fstest.MapFS{
+		"x/d":      &fstest.MapFile{Mode: fs.ModeDir},
+		"x/dd/a/d": &fstest.MapFile{Mode: fs.ModeDir},
+		"x/f":      &fstest.MapFile{Data: []byte("file")},
+		"x/ff/a/f": &fstest.MapFile{Data: []byte("file")},
+	}
+
+	p := Program{StdinIsTerminal: true}
+	p.Find = func(s string) fs.FS {
+		return must(fs.Sub(testfs, s))
+	}
+
+	p.Args = []string{`$find | keys[]`}
+	testRun(t, "", "error", &p)
+
+	p.Args = slices.Insert(p.Args, 0, "-find", "x")
+	testRun(t, "", "d/ dd/ dd/a/ dd/a/d/ f ff/ ff/a/ ff/a/f", &p)
+}
+
 func TestDry(t *testing.T) {
 	const q = `snapshot("\(.).json"; [.])`
 	p := Program{}