Add AutoInt option (#671)

This commit is contained in:
Shuhei Kitagawa 2025-02-25 14:44:58 +01:00 committed by GitHub
parent abc70836f5
commit 944206aba5
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
3 changed files with 61 additions and 0 deletions

View file

@ -38,6 +38,7 @@ type Encoder struct {
anchorNameMap map[string]struct{} anchorNameMap map[string]struct{}
anchorCallback func(*ast.AnchorNode, interface{}) error anchorCallback func(*ast.AnchorNode, interface{}) error
customMarshalerMap map[reflect.Type]func(interface{}) ([]byte, error) customMarshalerMap map[reflect.Type]func(interface{}) ([]byte, error)
autoInt bool
useLiteralStyleIfMultiline bool useLiteralStyleIfMultiline bool
commentMap map[*Path][]*Comment commentMap map[*Path][]*Comment
written bool written bool
@ -547,6 +548,9 @@ func (e *Encoder) encodeFloat(v float64, bitSize int) ast.Node {
} }
value := strconv.FormatFloat(v, 'g', -1, bitSize) value := strconv.FormatFloat(v, 'g', -1, bitSize)
if !strings.Contains(value, ".") && !strings.Contains(value, "e") { 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 // append x.0 suffix to keep float value context
value = fmt.Sprintf("%s.0", value) value = fmt.Sprintf("%s.0", value)
} }

View file

@ -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) { func TestEncoder_MultipleDocuments(t *testing.T) {
var buf bytes.Buffer var buf bytes.Buffer
enc := yaml.NewEncoder(&buf) enc := yaml.NewEncoder(&buf)

View file

@ -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. // CommentPosition type of the position for comment.
type CommentPosition int type CommentPosition int