Merge pull request #2726 from jandubois/lima-marshal

Refactor limayaml.Save and store.SaveYAML → limayaml.Marshal

Author: AkihiroSuda

cmd/limactl/validate.go

@@ -4,6 +4,7 @@ import (
 	"fmt"
 
 	"github.com/lima-vm/lima/cmd/limactl/guessarg"
+	"github.com/lima-vm/lima/pkg/limayaml"
 	"github.com/lima-vm/lima/pkg/store"
 	"github.com/spf13/cobra"
 
@@ -38,7 +39,7 @@ func validateAction(cmd *cobra.Command, args []string) error {
 		}
 		logrus.Infof("%q: OK", f)
 		if fill {
-			b, err := store.SaveYAML(y, len(args) > 1)
+			b, err := limayaml.Marshal(y, len(args) > 1)
 			if err != nil {
 				return err
 			}

pkg/limayaml/limayaml_test.go

@@ -21,7 +21,7 @@ const emptyYAML = "images: []\n"
 func TestEmptyYAML(t *testing.T) {
 	var y LimaYAML
 	t.Log(dumpJSON(t, y))
-	b, err := marshalYAML(y)
+	b, err := Marshal(&y, false)
 	assert.NilError(t, err)
 	assert.Equal(t, string(b), emptyYAML)
 }
@@ -32,12 +32,12 @@ func TestDefaultYAML(t *testing.T) {
 	bytes, err := os.ReadFile("default.yaml")
 	assert.NilError(t, err)
 	var y LimaYAML
-	err = unmarshalYAML(bytes, &y, "")
+	err = Unmarshal(bytes, &y, "")
 	assert.NilError(t, err)
 	y.Images = nil // remove default images
 	y.Mounts = nil // remove default mounts
 	t.Log(dumpJSON(t, y))
-	b, err := marshalYAML(y)
+	b, err := Marshal(&y, false)
 	assert.NilError(t, err)
 	assert.Equal(t, string(b), defaultYAML)
 }

pkg/limayaml/load.go

@@ -6,46 +6,18 @@ import (
 	"os"
 	"path/filepath"
 
-	"github.com/goccy/go-yaml"
 	"github.com/lima-vm/lima/pkg/store/dirnames"
 	"github.com/lima-vm/lima/pkg/store/filenames"
-	"github.com/lima-vm/lima/pkg/yqutil"
 	"github.com/sirupsen/logrus"
 )
 
-func unmarshalDisk(dst *Disk, b []byte) error {
-	var s string
-	if err := yaml.Unmarshal(b, &s); err == nil {
-		*dst = Disk{Name: s}
-		return nil
-	}
-	return yaml.Unmarshal(b, dst)
-}
-
-func unmarshalYAML(data []byte, v interface{}, comment string) error {
-	if err := yaml.UnmarshalWithOptions(data, v, yaml.DisallowDuplicateKey(), yaml.CustomUnmarshaler[Disk](unmarshalDisk)); err != nil {
-		return fmt.Errorf("failed to unmarshal YAML (%s): %w", comment, err)
-	}
-	// the go-yaml library doesn't catch all markup errors, unfortunately
-	// make sure to get a "second opinion", using the same library as "yq"
-	if err := yqutil.ValidateContent(data); err != nil {
-		return fmt.Errorf("failed to unmarshal YAML (%s): %w", comment, err)
-	}
-	if err := yaml.UnmarshalWithOptions(data, v, yaml.Strict(), yaml.CustomUnmarshaler[Disk](unmarshalDisk)); err != nil {
-		logrus.WithField("comment", comment).WithError(err).Warn("Non-strict YAML is deprecated and will be unsupported in a future version of Lima")
-		// Non-strict YAML is known to be used by Rancher Desktop:
-		// https://github.com/rancher-sandbox/rancher-desktop/blob/c7ea7508a0191634adf16f4675f64c73198e8d37/src/backend/lima.ts#L114-L117
-	}
-	return nil
-}
-
 // Load loads the yaml and fulfills unspecified fields with the default values.
 //
 // Load does not validate. Use Validate for validation.
 func Load(b []byte, filePath string) (*LimaYAML, error) {
 	var y, d, o LimaYAML
 
-	if err := unmarshalYAML(b, &y, fmt.Sprintf("main file %q", filePath)); err != nil {
+	if err := Unmarshal(b, &y, fmt.Sprintf("main file %q", filePath)); err != nil {
 		return nil, err
 	}
 	configDir, err := dirnames.LimaConfigDir()
@@ -57,7 +29,7 @@ func Load(b []byte, filePath string) (*LimaYAML, error) {
 	bytes, err := os.ReadFile(defaultPath)
 	if err == nil {
 		logrus.Debugf("Mixing %q into %q", defaultPath, filePath)
-		if err := unmarshalYAML(bytes, &d, fmt.Sprintf("default file %q", defaultPath)); err != nil {
+		if err := Unmarshal(bytes, &d, fmt.Sprintf("default file %q", defaultPath)); err != nil {
 			return nil, err
 		}
 	} else if !errors.Is(err, os.ErrNotExist) {
@@ -68,7 +40,7 @@ func Load(b []byte, filePath string) (*LimaYAML, error) {
 	bytes, err = os.ReadFile(overridePath)
 	if err == nil {
 		logrus.Debugf("Mixing %q into %q", overridePath, filePath)
-		if err := unmarshalYAML(bytes, &o, fmt.Sprintf("override file %q", overridePath)); err != nil {
+		if err := Unmarshal(bytes, &o, fmt.Sprintf("override file %q", overridePath)); err != nil {
 			return nil, err
 		}
 	} else if !errors.Is(err, os.ErrNotExist) {

pkg/limayaml/marshal.go

@@ -0,0 +1,62 @@
+package limayaml
+
+import (
+	"fmt"
+
+	"github.com/goccy/go-yaml"
+	"github.com/lima-vm/lima/pkg/yqutil"
+	"github.com/sirupsen/logrus"
+)
+
+func marshalString(s string) ([]byte, error) {
+	if s == "null" || s == "~" {
+		// work around go-yaml bugs
+		return []byte("\"" + s + "\""), nil
+	}
+	return yaml.Marshal(s)
+}
+
+const (
+	documentStart = "---\n"
+	documentEnd   = "...\n"
+)
+
+// Marshal the struct as a YAML document, optionally as a stream.
+func Marshal(y *LimaYAML, stream bool) ([]byte, error) {
+	options := []yaml.EncodeOption{yaml.CustomMarshaler[string](marshalString)}
+	b, err := yaml.MarshalWithOptions(y, options...)
+	if err != nil {
+		return nil, err
+	}
+	if stream {
+		doc := documentStart + string(b) + documentEnd
+		b = []byte(doc)
+	}
+	return b, nil
+}
+
+func unmarshalDisk(dst *Disk, b []byte) error {
+	var s string
+	if err := yaml.Unmarshal(b, &s); err == nil {
+		*dst = Disk{Name: s}
+		return nil
+	}
+	return yaml.Unmarshal(b, dst)
+}
+
+func Unmarshal(data []byte, v interface{}, comment string) error {
+	if err := yaml.UnmarshalWithOptions(data, v, yaml.DisallowDuplicateKey(), yaml.CustomUnmarshaler[Disk](unmarshalDisk)); err != nil {
+		return fmt.Errorf("failed to unmarshal YAML (%s): %w", comment, err)
+	}
+	// the go-yaml library doesn't catch all markup errors, unfortunately
+	// make sure to get a "second opinion", using the same library as "yq"
+	if err := yqutil.ValidateContent(data); err != nil {
+		return fmt.Errorf("failed to unmarshal YAML (%s): %w", comment, err)
+	}
+	if err := yaml.UnmarshalWithOptions(data, v, yaml.Strict(), yaml.CustomUnmarshaler[Disk](unmarshalDisk)); err != nil {
+		logrus.WithField("comment", comment).WithError(err).Warn("Non-strict YAML is deprecated and will be unsupported in a future version of Lima")
+		// Non-strict YAML is known to be used by Rancher Desktop:
+		// https://github.com/rancher-sandbox/rancher-desktop/blob/c7ea7508a0191634adf16f4675f64c73198e8d37/src/backend/lima.ts#L114-L117
+	}
+	return nil
+}

pkg/limayaml/save_test.go renamed to pkg/limayaml/marshal_test.go

@@ -7,29 +7,31 @@ import (
 	"gotest.tools/v3/assert"
 )
 
-func TestSaveEmpty(t *testing.T) {
-	_, err := Save(&LimaYAML{})
+func TestMarshalEmpty(t *testing.T) {
+	_, err := Marshal(&LimaYAML{}, false)
 	assert.NilError(t, err)
 }
 
-func TestSaveTilde(t *testing.T) {
+func TestMarshalTilde(t *testing.T) {
 	y := LimaYAML{
 		Mounts: []Mount{
 			{Location: "~", Writable: ptr.Of(false)},
 			{Location: "/tmp/lima", Writable: ptr.Of(true)},
 			{Location: "null"},
 		},
 	}
-	b, err := Save(&y)
+	b, err := Marshal(&y, true)
 	assert.NilError(t, err)
 	// yaml will load ~ (or null) as null
 	// make sure that it is always quoted
-	assert.Equal(t, string(b), `images: []
+	assert.Equal(t, string(b), `---
+images: []
 mounts:
 - location: "~"
   writable: false
 - location: /tmp/lima
   writable: true
 - location: "null"
+...
 `)
 }

pkg/limayaml/save.go

@@ -133,20 +133,3 @@ func LoadYAMLByFilePath(filePath string) (*limayaml.LimaYAML, error) {
 	}
 	return y, nil
 }
-
-const documentStart = "---\n"
-
-const documentEnd = "...\n"
-
-// SaveYAML saves the yaml, optionally as a stream.
-func SaveYAML(y *limayaml.LimaYAML, stream bool) ([]byte, error) {
-	b, err := limayaml.Save(y)
-	if err != nil {
-		return nil, err
-	}
-	if stream {
-		doc := documentStart + string(b) + documentEnd
-		b = []byte(doc)
-	}
-	return b, nil
-}