Skip to content

Commit 944206a

Browse files
authored
Add AutoInt option (#671)
1 parent abc7083 commit 944206a

File tree

3 files changed

+61
-0
lines changed

3 files changed

+61
-0
lines changed

‎encode.go‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ type Encoder struct {
3838
anchorNameMap map[string]struct{}
3939
anchorCallback func(*ast.AnchorNode, interface{}) error
4040
customMarshalerMap map[reflect.Type]func(interface{}) ([]byte, error)
41+
autoInt bool
4142
useLiteralStyleIfMultiline bool
4243
commentMap map[*Path][]*Comment
4344
written bool
@@ -547,6 +548,9 @@ func (e *Encoder) encodeFloat(v float64, bitSize int) ast.Node {
547548
}
548549
value := strconv.FormatFloat(v, 'g', -1, bitSize)
549550
if !strings.Contains(value, ".") && !strings.Contains(value, "e") {
551+
if e.autoInt {
552+
return ast.Integer(token.New(value, value, e.pos(e.column)))
553+
}
550554
// append x.0 suffix to keep float value context
551555
value = fmt.Sprintf("%s.0", value)
552556
}

‎encode_test.go‎

Lines changed: 48 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1283,6 +1283,54 @@ func TestEncoder_CustomMarshaler(t *testing.T) {
12831283
})
12841284
}
12851285

1286+
func TestEncoder_AutoInt(t *testing.T) {
1287+
for _, test := range []struct {
1288+
desc string
1289+
input any
1290+
expected string
1291+
}{
1292+
{
1293+
desc: "int-convertible float64",
1294+
input: map[string]float64{
1295+
"key": 1.0,
1296+
},
1297+
expected: "key: 1\n",
1298+
},
1299+
{
1300+
desc: "non int-convertible float64",
1301+
input: map[string]float64{
1302+
"key": 1.1,
1303+
},
1304+
expected: "key: 1.1\n",
1305+
},
1306+
{
1307+
desc: "int-convertible float32",
1308+
input: map[string]float32{
1309+
"key": 1.0,
1310+
},
1311+
expected: "key: 1\n",
1312+
},
1313+
{
1314+
desc: "non int-convertible float32",
1315+
input: map[string]float32{
1316+
"key": 1.1,
1317+
},
1318+
expected: "key: 1.1\n",
1319+
},
1320+
} {
1321+
t.Run(test.desc, func(t *testing.T) {
1322+
var buf bytes.Buffer
1323+
enc := yaml.NewEncoder(&buf, yaml.AutoInt())
1324+
if err := enc.Encode(test.input); err != nil {
1325+
t.Fatalf("failed to encode: %s", err)
1326+
}
1327+
if actual := buf.String(); actual != test.expected {
1328+
t.Errorf("expect:\n%s\nactual\n%s\n", test.expected, actual)
1329+
}
1330+
})
1331+
}
1332+
}
1333+
12861334
func TestEncoder_MultipleDocuments(t *testing.T) {
12871335
var buf bytes.Buffer
12881336
enc := yaml.NewEncoder(&buf)

‎option.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -206,6 +206,15 @@ func CustomMarshaler[T any](marshaler func(T) ([]byte, error)) EncodeOption {
206206
}
207207
}
208208

209+
// AutoInt automatically converts floating-point numbers to integers when the fractional part is zero.
210+
// For example, a value of 1.0 will be encoded as 1.
211+
func AutoInt() EncodeOption {
212+
return func(e *Encoder) error {
213+
e.autoInt = true
214+
return nil
215+
}
216+
}
217+
209218
// CommentPosition type of the position for comment.
210219
type CommentPosition int
211220

0 commit comments

Comments
 (0)