timplit.go

package timplit

import (
	"bufio"
	"bytes"
	"encoding/json"
	"io"
	"text/template"

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

var (
	newline   = []byte("\n")
	separator = []byte("---")
)

// Data is the JSON/YAML will be unmarshaled into this type, then applied to
// the template.
type Data interface{}

type Timplit struct {
	Template          io.Reader
	Data              io.Reader
	Out               io.Writer
	Err               io.Writer
	IgnoreParseErrors bool
}

// Execute applies the data to the template.
func Execute(templateText []byte, d *Data) ([]byte, error) {
	name := "timplit"

	// Create template
	t, err := template.New(name).Funcs(sprig.FuncMap()).Parse(string(templateText))
	if err != nil {
		return nil, err
	}

	// Execute template
	out := bytes.Buffer{}
	if err := t.ExecuteTemplate(&out, name, *d); err != nil {
		return nil, err
	}

	return out.Bytes(), nil
}

// ExecuteYAML applies the raw YAML text (a single document) to the template.
func ExecuteYAML(templateText []byte, yamlText []byte) ([]byte, error) {
	var data Data
	if err := yaml.Unmarshal(yamlText, &data); err != nil {
		return nil, err
	}
	return Execute(templateText, &data)
}

// ExecuteJSON applies the raw JSON text to the template.
func ExecuteJSON(templateText []byte, jsonText []byte) ([]byte, error) {
	var data Data
	if err := json.Unmarshal(jsonText, &data); err != nil {
		return nil, err
	}
	return Execute(templateText, &data)
}

// JSON reads a single JSON document from an io.Reader (usually stdin) until
// EOF, executes the template, and writes the result to an io.Writer (usually
// stdout). The template is always read first in full, and blocks until EOF is
// reached.
//
// If t.IgnoreParseErrors is true, then errors will be printed but execution
// will continue. Otherwise, the first error encountered is returned and
// subsequent lines will not be read.
func (t *Timplit) JSON() error {
	rawTemplate, err := io.ReadAll(t.Template)
	if err != nil {
		return err
	}
	rawJSON, err := io.ReadAll(t.Data)
	if err != nil {
		return err
	}
	out, err := ExecuteJSON(rawTemplate, rawJSON)
	if err != nil {
		return err
	}
	t.Out.Write(out)
	return nil
}

// MultipartJSON reads one JSON document per line from an io.Reader (usually
// stdin), executes the template for each line, and writes the result to an
// io.Writer (usually stdout). The template is always read first in full, and
// blocks until EOF is reached.
//
// If t.IgnoreParseErrors is true, then errors will be printed but execution
// will continue. Otherwise, the first error encountered is returned and
// subsequent lines will not be read.
func (t *Timplit) MultipartJSON() error {
	rawTemplate, err := io.ReadAll(t.Template)
	if err != nil {
		return err
	}

	// Read each line, which should be a complete JSON document
	scanner := bufio.NewScanner(t.Data)
	for scanner.Scan() {
		line := scanner.Bytes()
		if len(line) == 0 {
			// Skip empty lines
			continue
		}
		out, err := ExecuteJSON(rawTemplate, line)
		if e := t.handleError(err); e != nil {
			return e
		}
		t.Out.Write(out)
		t.Out.Write([]byte("\n"))
	}
	return nil
}

// YAML executes the template for each YAML document (separated by "---") it
// reads from an io.Reader (usually stdin) and writes the result to an
// io.Writer (usually stdout). The template is always read first in full, and
// blocks until EOF is reached.
//
// If t.IgnoreParseErrors is true, then errors will be printed but execution
// will continue. Otherwise, the first error encountered is returned and
// subsequent lines will not be read.
func (t *Timplit) YAML() error {
	rawTemplate, err := io.ReadAll(t.Template)
	if err != nil {
		return err
	}

	// The canonical (haha) way of doing this is with a yaml.Decoder, however
	// this does not allow continuing to the next document after a parse error in
	// another document.
	doc := bytes.Buffer{}
	s := bufio.NewScanner(t.Data)
	for s.Scan() {
		line := s.Bytes()
		if suffix, found := bytes.CutPrefix(line, separator); found {
			if err := t.executeYAML(rawTemplate, &doc); err != nil {
				return err
			}
			line = suffix
		}
		if !(doc.Len() == 0 && len(line) == 0) {
			doc.Write(line)
			doc.Write(newline)
		}
	}
	if err := t.executeYAML(rawTemplate, &doc); err != nil {
		return err
	}
	return nil
}

// executeYAML takes a single YAML document as an io.Reader and applies it to
// the template if it is not empty. Results are written to the t.Out Writer.
//
// If t.IgnoreParseErrors is true, then errors will be printed but execution
// will continue. Otherwise, the first error encountered is returned and
// subsequent lines will not be read.
func (t *Timplit) executeYAML(template []byte, doc io.Reader) error {
	data, err := io.ReadAll(doc)
	if e := t.handleError(err); e != nil {
		return e
	}
	if len(data) > 0 {
		out, err := ExecuteYAML(template, data)
		if e := t.handleError(err); e != nil {
			return e
		}
		t.Out.Write(out)
		t.Out.Write(newline)
	}
	return nil
}

// handleError will write the error message to t.Err and return nil if
// t.IgnoreParseErrors is set. Otherwise, it returns the original error.
func (t *Timplit) handleError(err error) error {
	if err != nil && t.IgnoreParseErrors {
		t.Err.Write([]byte("Error: " + err.Error()))
		return nil
	}
	return err
}