diff --git a/cmd/run.go b/cmd/run.go index e197015..272779b 100644 --- a/cmd/run.go +++ b/cmd/run.go @@ -41,12 +41,16 @@ var runCmd = &cobra.Command{ builder.WriteString("\n\n") // 2. preset files - rawFiles, ok := cfg.Preset[delineate] + var patterns []string + + raw, ok := cfg.Preset[delineate] if !ok { return fmt.Errorf("preset not found: %s", delineate) } - files, err := utils.ExpandGlob(rawFiles) + patterns = append(patterns, raw...) + + files, err := utils.ExpandGlob(patterns) if err != nil { return err } diff --git a/internal/config/config.go b/internal/config/config.go index f756aef..99faecd 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -1,11 +1,36 @@ package config import ( + "fmt" "os" "gopkg.in/yaml.v3" ) +type StringOrSlice []string + +func (s *StringOrSlice) UnmarshalYAML(value *yaml.Node) error { + switch value.Kind { + + // "preset: *" + case yaml.ScalarNode: + *s = []string{value.Value} + return nil + + // "preset: [a, b]" + case yaml.SequenceNode: + var out []string + for _, v := range value.Content { + out = append(out, v.Value) + } + *s = out + return nil + + default: + return fmt.Errorf("invalid preset format") + } +} + type Config struct { Config struct { PathMode string `yaml:"path_mode"` @@ -15,7 +40,7 @@ type Config struct { Preset string `yaml:"preset"` } `yaml:"prompts"` - Preset map[string][]string `yaml:"preset"` + Preset map[string]StringOrSlice `yaml:"preset"` } func Load() (*Config, error) { @@ -26,5 +51,14 @@ func Load() (*Config, error) { var cfg Config err = yaml.Unmarshal(data, &cfg) - return &cfg, err + if err != nil { + return nil, err + } + + // DEFAULT PROMPT TEMPLATE + if cfg.Prompts.Preset == "" { + cfg.Prompts.Preset = "The following code is from file {{path}}:\n---\n{{content}}" + } + + return &cfg, nil } \ No newline at end of file diff --git a/internal/utils/glob.go b/internal/utils/glob.go index 02d2618..920f6f5 100644 --- a/internal/utils/glob.go +++ b/internal/utils/glob.go @@ -3,30 +3,76 @@ package utils import ( "os" "path/filepath" + "strings" ) func ExpandGlob(patterns []string) ([]string, error) { var result []string + seen := make(map[string]struct{}) - for _, p := range patterns { - matches, err := filepath.Glob(p) + for _, pattern := range patterns { + + if pattern == "*" { + filepath.Walk(".", func(path string, info os.FileInfo, err error) error { + if err != nil { + return nil + } + if info.IsDir() { + return nil + } + add(path, seen, &result) + return nil + }) + continue + } + + isRecursive := strings.Contains(pattern, "**") + + if isRecursive { + root := strings.Split(pattern, "**")[0] + if root == "" { + root = "." + } + + filepath.Walk(root, func(path string, info os.FileInfo, err error) error { + if err != nil || info.IsDir() { + return nil + } + + matched, _ := filepath.Match( + strings.Replace(pattern, "**", "*", 1), + path, + ) + + if matched { + add(path, seen, &result) + } + return nil + }) + continue + } + + matches, err := filepath.Glob(pattern) if err != nil { return nil, err } for _, m := range matches { info, err := os.Stat(m) - if err != nil { + if err != nil || info.IsDir() { continue } - - if info.IsDir() { - continue - } - - result = append(result, m) + add(m, seen, &result) } } return result, nil +} + +func add(path string, seen map[string]struct{}, result *[]string) { + if _, ok := seen[path]; ok { + return + } + seen[path] = struct{}{} + *result = append(*result, path) } \ No newline at end of file