From 944206aba524110bfefdf0c53d864ce45257651c Mon Sep 17 00:00:00 2001
From: Shuhei Kitagawa <shuheiktgw@users.noreply.github.com>
Date: Tue, 25 Feb 2025 14:44:58 +0100
Subject: [PATCH] Add AutoInt option (#671)

---
 encode.go      |  4 ++++
 encode_test.go | 48 ++++++++++++++++++++++++++++++++++++++++++++++++
 option.go      |  9 +++++++++
 3 files changed, 61 insertions(+)

diff --git a/encode.go b/encode.go
index cc2961f..0a83017 100644
--- a/encode.go
+++ b/encode.go
@@ -38,6 +38,7 @@ type Encoder struct {
 	anchorNameMap              map[string]struct{}
 	anchorCallback             func(*ast.AnchorNode, interface{}) error
 	customMarshalerMap         map[reflect.Type]func(interface{}) ([]byte, error)
+	autoInt                    bool
 	useLiteralStyleIfMultiline bool
 	commentMap                 map[*Path][]*Comment
 	written                    bool
@@ -547,6 +548,9 @@ func (e *Encoder) encodeFloat(v float64, bitSize int) ast.Node {
 	}
 	value := strconv.FormatFloat(v, 'g', -1, bitSize)
 	if !strings.Contains(value, ".") && !strings.Contains(value, "e") {
+		if e.autoInt {
+			return ast.Integer(token.New(value, value, e.pos(e.column)))
+		}
 		// append x.0 suffix to keep float value context
 		value = fmt.Sprintf("%s.0", value)
 	}
diff --git a/encode_test.go b/encode_test.go
index be377c4..ef507c9 100644
--- a/encode_test.go
+++ b/encode_test.go
@@ -1283,6 +1283,54 @@ func TestEncoder_CustomMarshaler(t *testing.T) {
 	})
 }
 
+func TestEncoder_AutoInt(t *testing.T) {
+	for _, test := range []struct {
+		desc     string
+		input    any
+		expected string
+	}{
+		{
+			desc: "int-convertible float64",
+			input: map[string]float64{
+				"key": 1.0,
+			},
+			expected: "key: 1\n",
+		},
+		{
+			desc: "non int-convertible float64",
+			input: map[string]float64{
+				"key": 1.1,
+			},
+			expected: "key: 1.1\n",
+		},
+		{
+			desc: "int-convertible float32",
+			input: map[string]float32{
+				"key": 1.0,
+			},
+			expected: "key: 1\n",
+		},
+		{
+			desc: "non int-convertible float32",
+			input: map[string]float32{
+				"key": 1.1,
+			},
+			expected: "key: 1.1\n",
+		},
+	} {
+		t.Run(test.desc, func(t *testing.T) {
+			var buf bytes.Buffer
+			enc := yaml.NewEncoder(&buf, yaml.AutoInt())
+			if err := enc.Encode(test.input); err != nil {
+				t.Fatalf("failed to encode: %s", err)
+			}
+			if actual := buf.String(); actual != test.expected {
+				t.Errorf("expect:\n%s\nactual\n%s\n", test.expected, actual)
+			}
+		})
+	}
+}
+
 func TestEncoder_MultipleDocuments(t *testing.T) {
 	var buf bytes.Buffer
 	enc := yaml.NewEncoder(&buf)
diff --git a/option.go b/option.go
index a3be217..7dcd2ba 100644
--- a/option.go
+++ b/option.go
@@ -206,6 +206,15 @@ func CustomMarshaler[T any](marshaler func(T) ([]byte, error)) EncodeOption {
 	}
 }
 
+// AutoInt automatically converts floating-point numbers to integers when the fractional part is zero.
+// For example, a value of 1.0 will be encoded as 1.
+func AutoInt() EncodeOption {
+	return func(e *Encoder) error {
+		e.autoInt = true
+		return nil
+	}
+}
+
 // CommentPosition type of the position for comment.
 type CommentPosition int