more testing

cmd/disable/tests/disable-a-key-already-disabled/stderr.golden

@@ -1,3 +1,2 @@
 ---- exec command line 0: [disable KEY_B]
 WARNING: The key [KEY_B] is already disabled
----- done command line 0: [disable KEY_B]

cmd/disable/tests/disable-a-key-already-disabled/stdout.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [disable KEY_B]
----- done command line 0: [disable KEY_B]

cmd/disable/tests/disable-a-key/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [disable KEY_B]
----- done command line 0: [disable KEY_B]

cmd/disable/tests/disable-a-key/stdout.golden

@@ -1,3 +1,2 @@
 ---- exec command line 0: [disable KEY_B]
 Key [KEY_B] was successfully disabled
----- done command line 0: [disable KEY_B]

cmd/disable/tests/invalid-key/stderr.golden

@@ -1,4 +1,3 @@
 ---- exec command line 0: [disable NONEXISTING_KEY]
 Error: Could not find KEY [NONEXISTING_KEY]
 Run 'dottie disable --help' for usage.
----- done command line 0: [disable NONEXISTING_KEY]

cmd/disable/tests/invalid-key/stdout.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [disable NONEXISTING_KEY]
----- done command line 0: [disable NONEXISTING_KEY]

cmd/enable/tests/enable-a-key-already-enabled/stderr.golden

@@ -1,3 +1,2 @@
 ---- exec command line 0: [enable KEY_B]
 WARNING: The key [KEY_B] is already enabled
----- done command line 0: [enable KEY_B]

cmd/enable/tests/enable-a-key-already-enabled/stdout.golden

@@ -1,3 +1,2 @@
 ---- exec command line 0: [enable KEY_B]
 Key [KEY_B] was successfully enabled
----- done command line 0: [enable KEY_B]

cmd/enable/tests/enable-a-key/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [enable KEY_B]
----- done command line 0: [enable KEY_B]

cmd/enable/tests/enable-a-key/stdout.golden

@@ -1,3 +1,2 @@
 ---- exec command line 0: [enable KEY_B]
 Key [KEY_B] was successfully enabled
----- done command line 0: [enable KEY_B]

cmd/enable/tests/invalid-key/stderr.golden

@@ -1,4 +1,3 @@
 ---- exec command line 0: [enable NONEXISTING_KEY]
 Error: Could not find KEY [NONEXISTING_KEY]
 Run 'dottie enable --help' for usage.
----- done command line 0: [enable NONEXISTING_KEY]

cmd/enable/tests/invalid-key/stdout.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [enable NONEXISTING_KEY]
----- done command line 0: [enable NONEXISTING_KEY]

cmd/print/tests/empty/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [print --no-color]
----- done command line 0: [print --no-color]

cmd/print/tests/empty/stdout.golden

@@ -1,4 +1,3 @@
 ---- exec command line 0: [print --no-color]
 
 
----- done command line 0: [print --no-color]

cmd/print/tests/full/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [print --pretty --no-color]
----- done command line 0: [print --pretty --no-color]

cmd/print/tests/full/stdout.golden

@@ -30,4 +30,3 @@ GROUP_TWO_A="WORLD"
 
 SECRET=KEY                                                                      
 
----- done command line 0: [print --pretty --no-color]

cmd/print/tests/simple-pretty/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [print --no-color]
----- done command line 0: [print --no-color]

cmd/print/tests/simple-pretty/stdout.golden

@@ -3,4 +3,3 @@ KEY_A="I'm key A"
 KEY_B="I'm key B"
 KEY_C="I'm key C"
 
----- done command line 0: [print --no-color]

cmd/print/tests/simple/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [print --no-color]
----- done command line 0: [print --no-color]

cmd/print/tests/simple/stdout.golden

@@ -3,4 +3,3 @@ KEY_A="I'm key A"
 KEY_B="I'm key B"
 KEY_C="I'm key C"
 
----- done command line 0: [print --no-color]

cmd/print/tests/specific-group/stderr.golden

@@ -1,2 +1 @@
 ---- exec command line 0: [print --no-color --group my-first-group]
----- done command line 0: [print --no-color --group my-first-group]

cmd/print/tests/specific-group/stdout.golden

@@ -3,4 +3,3 @@ KEY_A="I'm key A"
 KEY_B="I'm key B"
 KEY_C="I'm key C"
 
----- done command line 0: [print --no-color --group my-first-group]

cmd/set/set.go

@@ -3,7 +3,6 @@ package set
 import (
 	"errors"
 	"fmt"
-	"os"
 	"strings"
 
 	"github.com/jippi/dottie/pkg"
@@ -87,7 +86,10 @@ func runE(cmd *cobra.Command, args []string) error {
 	// Loop arguments and place them
 	//
 
-	var allErrors error
+	var (
+		allErrors      error
+		stdout, stderr = tui.PrintersFromContext(cmd.Context())
+	)
 
 	for _, stringPair := range args {
 		pairSlice := strings.SplitN(stringPair, "=", 2)
@@ -106,7 +108,7 @@ func runE(cmd *cobra.Command, args []string) error {
 			Interpolated: value,
 			Enabled:      !shared.BoolFlag(cmd.Flags(), "disabled"),
 			Quote:        token.QuoteFromString(shared.StringFlag(cmd.Flags(), "quote-style")),
-			Comments:     ast.NewCommentsFromSlice(shared.StringSliceFlag(cmd.Flags(), "comments")),
+			Comments:     ast.NewCommentsFromSlice(shared.StringSliceFlag(cmd.Flags(), "comment")),
 		}
 
 		//
@@ -115,12 +117,12 @@ func runE(cmd *cobra.Command, args []string) error {
 
 		assignment, warnings, err := upserter.Upsert(assignment)
 		if warnings != nil {
-			tui.Theme.Warning.StderrPrinter().Println("WARNING:", warnings)
+			stderr.Color(tui.Warning).Println("WARNING:", warnings)
 		}
 
 		if err != nil {
 			z := validation.NewError(assignment, err)
-			fmt.Fprintln(os.Stderr, validation.Explain(document, z, z, false, true))
+			stderr.Color(tui.Neutral).Println(validation.Explain(document, z, z, false, true))
 
 			if shared.BoolWithInverseValue(cmd.Flags(), "validate") {
 				allErrors = multierr.Append(allErrors, err)
@@ -129,7 +131,7 @@ func runE(cmd *cobra.Command, args []string) error {
 			}
 		}
 
-		tui.Theme.Success.StderrPrinter().Printfln("Key [%s] was successfully upserted", key)
+		stdout.Color(tui.Success).Printfln("Key [%s] was successfully upserted", key)
 	}
 
 	if allErrors != nil {
@@ -144,7 +146,7 @@ func runE(cmd *cobra.Command, args []string) error {
 		return fmt.Errorf("failed to save file: %w", err)
 	}
 
-	tui.Theme.Success.StderrPrinter().Println("File was successfully saved")
+	stdout.Color(tui.Success).Println("File was successfully saved")
 
 	return nil
 }

cmd/set/set_test.go

@@ -0,0 +1,13 @@
+package set_test
+
+import (
+	"testing"
+
+	"github.com/jippi/dottie/pkg/test_helpers"
+)
+
+func TestCommand(t *testing.T) {
+	t.Parallel()
+
+	test_helpers.RunFileBasedCommandTests(t, 0, "set")
+}

cmd/set/tests/manipulate-empty.env

@@ -0,0 +1 @@
+SOME_KEY="SOME_VALUE"

cmd/set/tests/manipulate-empty.run

@@ -0,0 +1,9 @@
+SOME_KEY=SOME_VALUE
+ANOTHER_KEY=ANOTHER_VALUE --quote-style single
+SECOND_KEY="should be before ANOTHER_KEY" --before ANOTHER_KEY
+TRUE_SECOND_KEY="HA, I'm after SOME_KEY, so I'm before ANOTHER_KEY now" --after SOME_KEY
+SECOND_KEY="damn, I'm the third key now"
+SOME_KEY=ANOTHER_VALUE --comment "I'm a comment" --comment "I'm another comment"
+A_NUMBER=1 --comment "@dottie/validate number"
+NOT_A_NUMBER=abc --comment "@dottie/validate number"
+A_NUMBER=2

cmd/set/tests/manipulate-empty/env.golden

@@ -0,0 +1,10 @@
+# I'm a comment
+# I'm another comment
+SOME_KEY="ANOTHER_VALUE"
+
+TRUE_SECOND_KEY="HA, I'm after SOME_KEY, so I'm before ANOTHER_KEY now"
+SECOND_KEY="damn, I'm the third key now"
+ANOTHER_KEY='ANOTHER_VALUE'
+
+# @dottie/validate number
+A_NUMBER="2"

cmd/set/tests/manipulate-empty/stderr.golden

@@ -0,0 +1,14 @@
+---- exec command line 0: [set SOME_KEY=SOME_VALUE]
+---- exec command line 1: [set ANOTHER_KEY=ANOTHER_VALUE --quote-style single]
+---- exec command line 2: [set SECOND_KEY=should be before ANOTHER_KEY --before ANOTHER_KEY]
+---- exec command line 3: [set TRUE_SECOND_KEY=HA, I'm after SOME_KEY, so I'm before ANOTHER_KEY now --after SOME_KEY]
+---- exec command line 4: [set SECOND_KEY=damn, I'm the third key now]
+---- exec command line 5: [set SOME_KEY=ANOTHER_VALUE --comment I'm a comment --comment I'm another comment]
+---- exec command line 6: [set A_NUMBER=1 --comment @dottie/validate number]
+---- exec command line 7: [set NOT_A_NUMBER=abc --comment @dottie/validate number]
+NOT_A_NUMBER (-:2)                                 
+  * (number) The value [abc] is not a valid number.
+
+Error: validation error
+Run 'dottie set --help' for usage.
+---- exec command line 8: [set A_NUMBER=2]

cmd/set/tests/manipulate-empty/stdout.golden

@@ -0,0 +1,25 @@
+---- exec command line 0: [set SOME_KEY=SOME_VALUE]
+Key [SOME_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 1: [set ANOTHER_KEY=ANOTHER_VALUE --quote-style single]
+Key [ANOTHER_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 2: [set SECOND_KEY=should be before ANOTHER_KEY --before ANOTHER_KEY]
+Key [SECOND_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 3: [set TRUE_SECOND_KEY=HA, I'm after SOME_KEY, so I'm before ANOTHER_KEY now --after SOME_KEY]
+Key [TRUE_SECOND_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 4: [set SECOND_KEY=damn, I'm the third key now]
+Key [SECOND_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 5: [set SOME_KEY=ANOTHER_VALUE --comment I'm a comment --comment I'm another comment]
+Key [SOME_KEY] was successfully upserted
+File was successfully saved
+---- exec command line 6: [set A_NUMBER=1 --comment @dottie/validate number]
+Key [A_NUMBER] was successfully upserted
+File was successfully saved
+---- exec command line 7: [set NOT_A_NUMBER=abc --comment @dottie/validate number]
+---- exec command line 8: [set A_NUMBER=2]
+Key [A_NUMBER] was successfully upserted
+File was successfully saved

pkg/ast/document.go

@@ -295,3 +295,40 @@ func (document *Document) Initialize() {
 		}
 	}
 }
+
+func (document *Document) Replace(assignment *Assignment) error {
+	existing := document.Get(assignment.Name)
+	if existing == nil {
+		return fmt.Errorf("No KEY named [%s] exists in the document", assignment.Name)
+	}
+
+	if existing.Group != nil {
+		for idx, stmt := range existing.Group.Statements {
+			val, ok := stmt.(*Assignment)
+			if !ok {
+				continue
+			}
+
+			if val.Name == assignment.Name {
+				existing.Group.Statements[idx] = assignment
+
+				return nil
+			}
+		}
+	}
+
+	for idx, stmt := range document.Statements {
+		val, ok := stmt.(*Assignment)
+		if !ok {
+			continue
+		}
+
+		if val.Name == assignment.Name {
+			document.Statements[idx] = assignment
+
+			return nil
+		}
+	}
+
+	return fmt.Errorf("Could not find+replace KEY named [%s] in document", assignment.Name)
+}

pkg/ast/upsert/upsert.go

@@ -5,6 +5,9 @@ import (
 	"slices"
 
 	"github.com/jippi/dottie/pkg/ast"
+	"github.com/jippi/dottie/pkg/parser"
+	"github.com/jippi/dottie/pkg/render"
+	"github.com/jippi/dottie/pkg/scanner"
 	"github.com/jippi/dottie/pkg/validation"
 	"go.uber.org/multierr"
 )
@@ -79,9 +82,6 @@ func (u *Upserter) Upsert(input *ast.Assignment) (*ast.Assignment, error, error)
 		if err != nil {
 			return nil, nil, err
 		}
-
-		// Recalculate the index order of all Statements (for interpolation)
-		u.document.ReindexStatements()
 	}
 
 	// Replace comments on the assignment if the Setting is on
@@ -93,15 +93,32 @@ func (u *Upserter) Upsert(input *ast.Assignment) (*ast.Assignment, error, error)
 	assignment.Literal = input.Literal
 	assignment.Quote = input.Quote
 	assignment.Interpolated = input.Literal
+
+	var (
+		tempDoc       *ast.Document
+		err, warnings error
+	)
+
+	// Render and parse back the Statement to ensure annotations and such are properly handled
+	tempDoc, err = parser.New(scanner.New(render.NewFormatter().Statement(assignment).String()), "-").Parse()
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to parse assignment: %w", err)
+	}
+
+	assignment = tempDoc.Get(assignment.Name)
 	assignment.Initialize()
 
 	if _, ok := assignment.Dependencies[assignment.Name]; ok {
 		return nil, nil, fmt.Errorf("Key [%s] may not reference itself!", assignment.Name)
 	}
 
-	u.document.Initialize()
+	// Replace the Assignment in the document
+	//
+	// This is necessary since its a different pointer address after we rendered+parsed earlier
+	u.document.Replace(assignment)
 
-	var err, warnings error
+	// Reinitialize the document so all indices and such are correct
+	u.document.Initialize()
 
 	// Interpolate the Assignment if it is enabled
 	if assignment.Enabled {
@@ -128,18 +145,25 @@ func (u *Upserter) Upsert(input *ast.Assignment) (*ast.Assignment, error, error)
 }
 
 func (u *Upserter) createAndInsert(input *ast.Assignment) (*ast.Assignment, error) {
-	// Ensure the group exists (may return 'nil' if no group is required)
-	group := u.document.EnsureGroup(u.group)
-
 	// Create the new newAssignment
 	newAssignment := &ast.Assignment{
 		Comments: input.Comments,
 		Enabled:  input.Enabled,
-		Group:    group,
 		Literal:  input.Literal,
 		Name:     input.Name,
 	}
 
+	doc, err := parser.New(scanner.New(render.NewFormatter().Statement(newAssignment).String()), "-").Parse()
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse assignment: %w", err)
+	}
+
+	// Ensure the group exists (may return 'nil' if no group is required)
+	group := u.document.EnsureGroup(u.group)
+
+	newAssignment = doc.Get(newAssignment.Name)
+	newAssignment.Group = group
+
 	// Find the statement slice to operate on
 	statements := u.document.Statements
 	if newAssignment.Group != nil {

pkg/test_helpers/filebased_command_tests.go

@@ -145,9 +145,6 @@ func RunFileBasedCommandTests(t *testing.T, settings Setting, globalArgs ...stri
 
 				// Assert we got a Cobra command back
 				require.NotNil(t, out, "expected a return value")
-
-				stdout.WriteString(fmt.Sprintf("---- done command line %d: %+v\n", idx, args))
-				stderr.WriteString(fmt.Sprintf("---- done command line %d: %+v\n", idx, args))
 			}
 
 			// Assert stdout + stderr + modified env file is as expected