🐛 Fix map and seq templates being stripped

cmd/cmd.go

@@ -5,7 +5,6 @@ import (
 	"errors"
 	"fmt"
 	"github.com/clevyr/yampl/internal/config"
-	"github.com/clevyr/yampl/internal/node"
 	"github.com/clevyr/yampl/internal/visitor"
 	log "github.com/sirupsen/logrus"
 	"github.com/spf13/cobra"
@@ -164,7 +163,7 @@ func templateReader(conf config.Config, r io.Reader) (string, error) {
 			buf.Write([]byte("---\n"))
 		}
 
-		if err := node.Visit(v.Visit, &n); err != nil {
+		if err := v.Run(&n); err != nil {
 			return buf.String(), err
 		}
 

cmd/flag_values.go

@@ -2,7 +2,6 @@ package cmd
 
 import (
 	"errors"
-	"github.com/clevyr/yampl/internal/node"
 	"github.com/clevyr/yampl/internal/visitor"
 	"github.com/spf13/cobra"
 	"gopkg.in/yaml.v3"
@@ -46,7 +45,7 @@ func valueCompletion(cmd *cobra.Command, args []string, toComplete string) ([]st
 				continue
 			}
 
-			if err := node.Visit(v.Visit, &n); err != nil {
+			if err := v.Run(&n); err != nil {
 				continue
 			}
 		}

internal/visitor/find_args.go

@@ -59,7 +59,22 @@ type FindArgs struct {
 	matchMap map[string]MatchSlice
 }
 
-func (visitor *FindArgs) Visit(n *yaml.Node) error {
+func (visitor *FindArgs) Run(n *yaml.Node) error {
+	if len(n.Content) == 0 {
+		if err := visitor.FindArgs(n); err != nil {
+			return err
+		}
+	} else {
+		for _, node := range n.Content {
+			if err := visitor.Run(node); err != nil {
+				return err
+			}
+		}
+	}
+	return nil
+}
+
+func (visitor *FindArgs) FindArgs(n *yaml.Node) error {
 	if tmplSrc, _ := node.GetCommentTmpl(visitor.conf.Prefix, n); tmplSrc != "" {
 		tmpl, err := template.New("").
 			Funcs(template2.FuncMap()).

internal/visitor/template_comments.go

@@ -21,66 +21,103 @@ type TemplateComments struct {
 	conf config.Config
 }
 
-func (t TemplateComments) Visit(n *yaml.Node) error {
-	if tmplSrc, tmplTag := node.GetCommentTmpl(t.conf.Prefix, n); tmplSrc != "" {
-		t.conf.Log = t.conf.Log.WithFields(log.Fields{
-			"tmpl":    tmplSrc,
-			"filePos": fmt.Sprintf("%d:%d", n.Line, n.Column),
-			"from":    n.Value,
-		})
-
-		tmpl, err := template.New("").
-			Funcs(template2.FuncMap()).
-			Delims(t.conf.LeftDelim, t.conf.RightDelim).
-			Option("missingkey=error").
-			Parse(tmplSrc)
-		if err != nil {
-			if !t.conf.Fail {
-				t.conf.Log.WithError(err).Warn("skipping value due to template error")
-				return nil
+func (t TemplateComments) Run(n *yaml.Node) error {
+	if len(n.Content) == 0 {
+		// Node has no children. Template current node.
+		tmplSrc, tmplTag := node.GetCommentTmpl(t.conf.Prefix, n)
+		if tmplSrc != "" {
+			if err := t.Template(n, tmplSrc, tmplTag); err != nil {
+				if t.conf.Fail {
+					return err
+				} else {
+					t.conf.Log.WithError(err).Warn("skipping value due to template error")
+				}
 			}
-			return NodeErr{Err: err, Node: n}
 		}
+		return nil
+	}
 
-		if t.conf.Values != nil {
-			t.conf.Values["Value"] = n.Value
-		}
+	switch n.Kind {
+	case yaml.MappingNode:
+		for i := 0; i < len(n.Content); i += 2 {
+			// Attempt to fetch template from comments on the key.
+			key, val := n.Content[i], n.Content[i+1]
 
-		var buf bytes.Buffer
-		if err = tmpl.Execute(&buf, t.conf.Values); err != nil {
-			if !t.conf.Fail {
-				t.conf.Log.WithError(err).Warn("skipping value due to template error")
-				return nil
+			tmplSrc, tmplTag := node.GetCommentTmpl(t.conf.Prefix, key)
+			if tmplSrc != "" {
+				if err := t.Template(val, tmplSrc, tmplTag); err != nil {
+					if t.conf.Fail {
+						return err
+					} else {
+						t.conf.Log.WithError(err).Warn("skipping value due to template error")
+					}
+				} else {
+					// Current node was templated, do not need to traverse children
+					continue
+				}
+			}
+
+			// Key did not have comment, run again with value.
+			if err := t.Run(val); err != nil {
+				return err
+			}
+		}
+	default:
+		for _, n := range n.Content {
+			if err := t.Run(n); err != nil {
+				return err
 			}
-			return NodeErr{Err: err, Node: n}
 		}
+	}
+	return nil
+}
 
-		if buf.String() != n.Value {
-			t.conf.Log.WithField("to", buf.String()).Debug("updating value")
-			n.Style = 0
+func (t TemplateComments) Template(n *yaml.Node, tmplSrc string, tmplTag node.TmplTag) error {
+	t.conf.Log = t.conf.Log.WithFields(log.Fields{
+		"tmpl":    tmplSrc,
+		"filePos": fmt.Sprintf("%d:%d", n.Line, n.Column),
+		"from":    n.Value,
+	})
 
-			switch tmplTag {
-			case node.SeqTag, node.MapTag:
-				var tmpNode yaml.Node
+	tmpl, err := template.New("").
+		Funcs(template2.FuncMap()).
+		Delims(t.conf.LeftDelim, t.conf.RightDelim).
+		Option("missingkey=error").
+		Parse(tmplSrc)
+	if err != nil {
+		return NodeErr{Err: err, Node: n}
+	}
 
-				if err := yaml.Unmarshal(buf.Bytes(), &tmpNode); err != nil {
-					if !t.conf.Fail {
-						t.conf.Log.WithError(err).Warn("skipping value due to unmarshal error")
-						return nil
-					}
-					return NodeErr{Err: err, Node: n}
-				}
+	if t.conf.Values != nil {
+		t.conf.Values["Value"] = n.Value
+	}
 
-				content := tmpNode.Content[0]
-				n.Content = content.Content
-				n.Kind = content.Kind
-				n.Value = content.Value
-			default:
-				n.SetString(buf.String())
+	var buf bytes.Buffer
+	if err = tmpl.Execute(&buf, t.conf.Values); err != nil {
+		return NodeErr{Err: err, Node: n}
+	}
+
+	if buf.String() != n.Value {
+		t.conf.Log.WithField("to", buf.String()).Debug("updating value")
+		n.Style = 0
+
+		switch tmplTag {
+		case node.SeqTag, node.MapTag:
+			var tmpNode yaml.Node
+
+			if err := yaml.Unmarshal(buf.Bytes(), &tmpNode); err != nil {
+				return NodeErr{Err: err, Node: n}
 			}
 
-			n.Tag = tmplTag.ToYaml()
+			content := tmpNode.Content[0]
+			n.Content = content.Content
+			n.Kind = content.Kind
+			n.Value = content.Value
+		default:
+			n.SetString(buf.String())
 		}
+
+		n.Tag = tmplTag.ToYaml()
 	}
 	return nil
 }