spry/main.go

package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"html/template"
	"io"
	"io/fs"
	"os"
	"path/filepath"
	"strings"
	"time"

	"github.com/peterbourgon/mergemap"

	sprig "github.com/Masterminds/sprig/v3"
	yaml "gopkg.in/yaml.v3"
)

type StrSet map[string]struct{}

func (s *StrSet) Add(str string) {
	(*s)[str] = struct{}{}
}

func (s *StrSet) Has(str string) bool {
	if _, ok := (*s)[str]; ok {
		return true
	}
	return false
}

func (s *StrSet) AddHas(str string) bool {
	if ok := s.Has(str); ok {
		return ok
	}
	s.Add(str)
	return false
}

type valFileSlice []string

func (v *valFileSlice) String() string {
	return fmt.Sprintf("%v", *v)
}

func (v *valFileSlice) Set(value string) error {
	*v = append(*v, value)
	return nil
}

type valEnt struct {
	address []string
	v       string
}

type valSlice []valEnt

func (v *valSlice) String() string {
	return fmt.Sprintf("%v", *v)
}

func (v *valSlice) Set(value string) error {
	pieces := strings.Split(value, "=")
	if len(pieces) < 2 {
		return fmt.Errorf("Invalid value setting.")
	}
	*v = append(*v, valEnt{
		strings.Split(pieces[0], "."),
		pieces[1],
	})
	return nil
}

var flags struct {
	valuesFiles valFileSlice
	values      valSlice
}

func init() {
	flag.Var(&flags.values, "set", "")
	flag.Var(&flags.valuesFiles, "f", "")
	flag.Parse()
}

func checkErr(err error) {
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %+v", err)
		os.Exit(1)
	}
}

func expandFiles(files []string) (expFiles []string) {
	expFiles = []string{}
	fileSet := StrSet{}
	directories := []string{}

	// Ensure we don't process any provided values files as templates.
	for _, path := range flags.valuesFiles {
		abs, err := filepath.Abs(path)
		if err != nil {
			fileSet.Add(path)
			break
		}
		fileSet.Add(abs)
	}

	for _, path := range files {
		fileInfo, err := os.Stat(path)
		checkErr(err)

		if fileInfo.IsDir() {
			// When provided path is a directory, defer loading files from it.
			directories = append(directories, path)
		} else {
			abs, err := filepath.Abs(path)
			if err != nil {
				// On failure to find the absolute path, just use the provided path.
				abs = path
			}
			// Deduplicate files.
			if ok := fileSet.AddHas(abs); !ok {
				expFiles = append(expFiles, path)
			}
		}
		// Walk each provided directory to find templates.
		for _, path := range directories {
			filepath.WalkDir(path, func(path string, d fs.DirEntry, err error) error {
				if err != nil {
					fmt.Fprintf(os.Stderr, "WARN: %+v", err)
					return err
				}
				if d.IsDir() {
					return nil
				}
				abs, err := filepath.Abs(path)
				if err != nil {
					// On failure to find the absolute path, just use the provided path.
					abs = path
				}
				// Deduplicate files.
				if ok := fileSet.AddHas(abs); !ok {
					expFiles = append(expFiles, path)
				}
				return nil
			})
		}
	}
	return
}

func mergeValues() (result map[string]interface{}) {
	result = map[string]interface{}{}
	for _, vf := range flags.valuesFiles {
		vfContent, err := os.ReadFile(vf)
		checkErr(err)
		if strings.HasSuffix(vf, ".json") {
			err = json.Unmarshal(vfContent, &result)
			checkErr(err)
		} else {
			err = yaml.Unmarshal(vfContent, result)
			checkErr(err)
		}
	}
	for _, v := range flags.values {
		valueMap := map[string]interface{}{}
		subMap := valueMap
		for i, node := range v.address {
			if i == len(v.address)-1 {
				subMap[node] = v.v
				break
			}
			subMap[node] = map[string]interface{}{}
			subMap = subMap[node].(map[string]interface{})
		}
		mergemap.Merge(result, valueMap)
	}
	return
}

func main() {
	var err error
	tpl := template.New("base").Funcs(sprig.FuncMap())

	if len(flag.Args()) == 0 {
		// With no provided arguments read from STDIN if anything is immediately available through it.
		// otherwise print the usage and exit.
		ch := make(chan struct{})
		go func() {
			select {
			case <-ch:
				return
			case <-time.After(100 * time.Millisecond):
				flag.Usage()
				os.Exit(2)
			}
		}()
		var data []byte
		data, err = io.ReadAll(os.Stdin)
		// Tell the waiting routine that we got input on STDIN.
		ch <- struct{}{}
		checkErr(err)
		tpl, err = tpl.Parse(string(data))
		checkErr(err)
	} else if len(flag.Args()) == 1 && flag.Args()[0] == "-" {
		// If only a single non-flag argument is provided and it's a -, wait for STDIN indefinitely.
		var data []byte
		data, err = io.ReadAll(os.Stdin)
		checkErr(err)
		tpl, err = tpl.Parse(string(data))
		checkErr(err)
	} else {
		// Otherwise use the provided files/directories for templates to render.
		// fmt.Fprintf(os.Stderr, "%+v\n", expandFiles(flag.Args()))
		tpl, err = tpl.ParseFiles(expandFiles(flag.Args())...)
		checkErr(err)
	}
	// fmt.Fprintf(os.Stderr, "%+v\n", mergeValues())
	tpl.Execute(os.Stdout, mergeValues())
}