fix: fix fmt toml. Only allow table after table in root scope (#430)

* fix: fix fmt toml. Only allow table after table in root scope Signed-off-by: he1pa <18012015693@163.com> * fix ut Signed-off-by: he1pa <18012015693@163.com> * add ut Signed-off-by: he1pa <18012015693@163.com> * record enc.inTable statu Signed-off-by: he1pa <18012015693@163.com> * fix ut Signed-off-by: he1pa <18012015693@163.com> --------- Signed-off-by: he1pa <18012015693@163.com>
kcl-lang · Jan 8, 2025 · 744add4 · 744add4
1 parent 59c6847
commit 744add4
Show file tree

Hide file tree

Showing 2 changed files with 54 additions and 15 deletions.
diff --git a/pkg/3rdparty/toml/encode.go b/pkg/3rdparty/toml/encode.go
@@ -127,9 +127,11 @@ func Marshal(v any) ([]byte, error) {
 // NOTE: only exported keys are encoded due to the use of reflection. Unexported
 // keys are silently discarded.
 type Encoder struct {
-	Indent     string // string for a single indentation level; default is two spaces.
-	hasWritten bool   // written any output to w yet?
-	w          *bufio.Writer
+	Indent          string // string for a single indentation level; default is two spaces.
+	hasWritten      bool   // written any output to w yet?
+	hasWrittenTable bool   // table must at end of file
+	inTable         bool
+	w               *bufio.Writer
 }
 
 // NewEncoder create a new Encoder.
@@ -168,6 +170,7 @@ func (enc *Encoder) encode(key Key, rv reflect.Value) {
 	// If we can marshal the type to text, then we use that. This prevents the
 	// encoder for handling these types as generic structs (or whatever the
 	// underlying type of a TextMarshaler is).
+
 	switch {
 	case isMarshaler(rv):
 		enc.writeKeyValue(key, rv, false)
@@ -183,6 +186,18 @@ func (enc *Encoder) encode(key Key, rv reflect.Value) {
 	}
 
 	k := rv.Kind()
+
+	// https://toml.io/en/v1.0.0#table
+	// Under a table,  and until the next header or EOF, are the key/values of that table.
+	// Therefore, if a table has been defined and is not within the scope of this table, only other tables are allowed to appear
+	if !enc.inTable && enc.hasWrittenTable {
+		switch k {
+		case reflect.Map, reflect.Struct:
+		default:
+			encPanic(fmt.Errorf("unsupported to define '%s' after a table, ref: https://toml.io/en/v1.0.0#table", key))
+		}
+	}
+
 	switch k {
 	case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32,
 		reflect.Int64,
@@ -379,16 +394,19 @@ func (enc *Encoder) eArrayOfTables(key Key, rv reflect.Value) {
 }
 
 func (enc *Encoder) eTable(key Key, rv reflect.Value) {
+	enc.inTable = true
 	if len(key) == 1 {
 		// Output an extra newline between top-level tables.
 		// (The newline isn't written if nothing else has been written though.)
 		enc.newline()
 	}
 	if len(key) > 0 {
+		enc.hasWrittenTable = true
 		enc.wf("%s[%s]", enc.indentStr(key), key)
 		enc.newline()
 	}
 	enc.eMapOrStruct(key, rv, false)
+	enc.inTable = false
 }
 
 func (enc *Encoder) eMapOrStruct(key Key, rv reflect.Value, inline bool) {

diff --git a/pkg/tools/gen/gentoml_test.go b/pkg/tools/gen/gentoml_test.go
@@ -1,6 +1,7 @@
 package gen
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/goccy/go-yaml"
@@ -30,22 +31,24 @@ c_key = "value3"
 		{
 			name: "Nested MapSlice",
 			data: &yaml.MapSlice{
+				{Key: "outer_key2", Value: "outer_value2"},
 				{Key: "outer_key1", Value: yaml.MapSlice{
 					{Key: "inner_key1", Value: "inner_value1"},
 					{Key: "inner_key2", Value: "inner_value2"},
 				}},
-				{Key: "outer_key2", Value: "outer_value2"},
 			},
-			expectedTOML: `[outer_key1]
+			expectedTOML: `outer_key2 = "outer_value2"
+
+[outer_key1]
   inner_key1 = "inner_value1"
   inner_key2 = "inner_value2"
-outer_key2 = "outer_value2"
 `,
 			expectErr: nil,
 		},
 		{
 			name: "Nested MapSlice with Slice",
 			data: &yaml.MapSlice{
+				{Key: "simple_key", Value: "simple_value"},
 				{Key: "key_with_slices", Value: []yaml.MapSlice{
 					{
 						{Key: "inner_key1", Value: "value1"},
@@ -54,10 +57,9 @@ outer_key2 = "outer_value2"
 						{Key: "inner_key2", Value: "value2"},
 					},
 				}},
-				{Key: "simple_key", Value: "simple_value"},
 			},
-			expectedTOML: `key_with_slices = [{inner_key1 = "value1"}, {inner_key2 = "value2"}]
-simple_key = "simple_value"
+			expectedTOML: `simple_key = "simple_value"
+key_with_slices = [{inner_key1 = "value1"}, {inner_key2 = "value2"}]
 `,
 			expectErr: nil,
 		},
@@ -86,31 +88,50 @@ simple_key = "simple_value"
 			name: "Simple MapSlice",
 			data: &yaml.MapSlice{
 				{Key: "b_key", Value: "value1"},
+				{Key: "c_key", Value: "value3"},
 				{Key: "a_key", Value: map[string]string{
 					"a_a_key": "value2",
 				}},
-				{Key: "c_key", Value: "value3"},
 			},
 			expectedTOML: `b_key = "value1"
+c_key = "value3"
 
 [a_key]
   a_a_key = "value2"
-c_key = "value3"
 `,
 			expectErr: nil,
 		},
+
+		{
+			name: "Table Test",
+			data: &yaml.MapSlice{
+				{Key: "a_key", Value: map[string]string{
+					"a_a_key": "value2",
+				}},
+				{Key: "b_key", Value: "value1"},
+			},
+			expectedTOML: "",
+			expectErr:    fmt.Errorf("unsupported to define 'b_key' after a table, ref: https://toml.io/en/v1.0.0#table"),
+		},
 	}
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
 			tomlData, err := toml.Marshal(tt.data)
-			if err != tt.expectErr {
+			if err != nil && tt.expectErr != nil {
+				if err.Error() != tt.expectErr.Error() {
+					t.Fatalf("expected error: %v, got: %v", tt.expectErr, err)
+				} else {
+					return
+				}
+			} else if err == nil && tt.expectErr == nil {
+				if got := string(tomlData); got != tt.expectedTOML {
+					t.Errorf("expected:\n%s\ngot:\n%s", tt.expectedTOML, got)
+				}
+			} else {
 				t.Fatalf("expected error: %v, got: %v", tt.expectErr, err)
 			}
 
-			if got := string(tomlData); got != tt.expectedTOML {
-				t.Errorf("expected:\n%s\ngot:\n%s", tt.expectedTOML, got)
-			}
 		})
 	}
 }