Repository: ogen-go/ogen
Pull request #525 (2022-08-08 08:56:02)

fix(jsonschema): marshal `Enum` and `Num` correctly
For gotd/botapi#217.


Patch:
diff --git a/jsonschema/enum.go b/jsonschema/enum.go
index 27f520a6..2a7adb19 100644
--- a/jsonschema/enum.go
+++ b/jsonschema/enum.go
@@ -2,6 +2,7 @@ package jsonschema
 
 import (
 	"encoding/json"
+	"reflect"
 
 	"github.com/go-faster/errors"
 	yaml "github.com/go-faster/yamlx"
@@ -18,15 +19,23 @@ func (n Enum) MarshalYAML() (interface{}, error) {
 		if err != nil {
 			return nil, err
 		}
-		vals = append(vals, &node)
+		vals = append(vals, node)
 	}
-	return vals, nil
+	return &yaml.Node{
+		Kind:    yaml.SequenceNode,
+		Tag:     "!!seq",
+		Content: vals,
+	}, nil
 }
 
 // UnmarshalYAML implements yaml.Unmarshaler.
 func (n *Enum) UnmarshalYAML(node *yaml.Node) error {
-	if node.Tag != "!!seq" {
-		return errors.Errorf("unexpected tag %s", node.Tag)
+	if node.Kind != yaml.SequenceNode {
+		return &yaml.UnmarshalError{
+			Node: node,
+			Type: reflect.TypeOf(n),
+			Err:  errors.Errorf("cannot unmarshal %s into %T", node.ShortTag(), n),
+		}
 	}
 	*n = (*n)[:0]
 	for _, val := range node.Content {
diff --git a/jsonschema/num.go b/jsonschema/num.go
index e50a5e46..50b8b961 100644
--- a/jsonschema/num.go
+++ b/jsonschema/num.go
@@ -2,6 +2,7 @@ package jsonschema
 
 import (
 	"encoding/json"
+	"reflect"
 
 	"github.com/go-faster/errors"
 	"github.com/go-faster/jx"
@@ -11,10 +12,27 @@ import (
 // Num represents JSON number.
 type Num json.RawMessage
 
+// MarshalYAML implements yaml.Marshaler.
+func (n Num) MarshalYAML() (interface{}, error) {
+	tag := "!!float"
+	if jx.Num(n).IsInt() {
+		tag = "!!int"
+	}
+	return yaml.Node{
+		Kind:  yaml.ScalarNode,
+		Tag:   tag,
+		Value: string(n),
+	}, nil
+}
+
 // UnmarshalYAML implements yaml.Unmarshaler.
 func (n *Num) UnmarshalYAML(node *yaml.Node) error {
-	if t := node.Tag; t != "!!int" && t != "!!float" {
-		return errors.Errorf("unexpected tag %s", t)
+	if t := node.Tag; node.Kind != yaml.ScalarNode || (t != "!!int" && t != "!!float") {
+		return &yaml.UnmarshalError{
+			Node: node,
+			Type: reflect.TypeOf(n),
+			Err:  errors.Errorf("cannot unmarshal %s into %T", node.ShortTag(), n),
+		}
 	}
 
 	val, err := convertYAMLtoRawJSON(node)
@@ -30,3 +48,21 @@ func (n *Num) UnmarshalYAML(node *yaml.Node) error {
 
 	return nil
 }
+
+// MarshalJSON implements json.Marshaler.
+func (n Num) MarshalJSON() ([]byte, error) {
+	return json.RawMessage(n).MarshalJSON()
+}
+
+// UnmarshalJSON implements json.Unmarshaler.
+func (n *Num) UnmarshalJSON(data []byte) error {
+	num, err := jx.DecodeBytes(data).Num()
+	if err != nil {
+		return err
+	}
+	if num.Str() {
+		return errors.Errorf("unexpected string %s", num)
+	}
+	*n = Num(num)
+	return nil
+}
diff --git a/jsonschema/raw_value.go b/jsonschema/raw_value.go
index 5a1b770b..eda8cd47 100644
--- a/jsonschema/raw_value.go
+++ b/jsonschema/raw_value.go
@@ -41,9 +41,15 @@ func (n *RawValue) UnmarshalJSON(b []byte) error {
 	return nil
 }
 
-func convertJSONToRawYAML(raw json.RawMessage) (node yaml.Node, err error) {
-	err = yaml.Unmarshal(raw, &node)
-	return node, err
+func convertJSONToRawYAML(raw json.RawMessage) (*yaml.Node, error) {
+	var node yaml.Node
+	if err := yaml.Unmarshal(raw, &node); err != nil {
+		return nil, err
+	}
+	if node.Kind == yaml.DocumentNode {
+		return node.Content[0], nil
+	}
+	return &node, nil
 }
 
 func convertYAMLtoRawJSON(node *yaml.Node) (json.RawMessage, error) {


Test patch:
diff --git a/jsonschema/enum_test.go b/jsonschema/enum_test.go
new file mode 100644
index 00000000..1d7ce78d
--- /dev/null
+++ b/jsonschema/enum_test.go
@@ -0,0 +1,39 @@
+package jsonschema
+
+import (
+	"encoding/json"
+	"fmt"
+	"testing"
+)
+
+func TestEnum(t *testing.T) {
+	create := func() interface{} {
+		return &Enum{}
+	}
+	enum := func(s ...string) (r Enum) {
+		for _, v := range s {
+			r = append(r, json.RawMessage(v))
+		}
+		return r
+	}
+
+	tests := []struct {
+		data    string
+		value   Enum
+		wantErr bool
+	}{
+		{`[0, 1]`, enum(`0`, `1`), false},
+		{`[0, "string"]`, enum(`0`, `"string"`), false},
+		{`[{}, []]`, enum(`{}`, `[]`), false},
+		// Invalid YAML.
+		{`"`, nil, true},
+		{`[`, nil, true},
+		// Invalid type.
+		{`{}`, nil, true},
+		{`"100"`, nil, true},
+	}
+	for i, tt := range tests {
+		tt := tt
+		t.Run(fmt.Sprintf("Test%d", i+1), testCustomEncodings(create, tt.data, tt.wantErr))
+	}
+}
diff --git a/jsonschema/num_test.go b/jsonschema/num_test.go
index 6c4f200a..4bb082ff 100644
--- a/jsonschema/num_test.go
+++ b/jsonschema/num_test.go
@@ -3,20 +3,22 @@ package jsonschema
 import (
 	"fmt"
 	"testing"
-
-	yaml "github.com/go-faster/yamlx"
-	"github.com/stretchr/testify/require"
 )
 
-func TestNum_UnmarshalYAML(t *testing.T) {
+func TestNum(t *testing.T) {
+	create := func() interface{} {
+		return &Num{}
+	}
+
 	tests := []struct {
 		data    string
 		value   Num
 		wantErr bool
 	}{
 		{`0`, Num(`0`), false},
-		{`1e1`, Num(`10`), false},
-		{`0x0a`, Num(`10`), false},
+		{`-1`, Num(`-1`), false},
+		{`1.1`, Num(`1.1`), false},
+		{`0.01e1`, Num(`0.1`), false},
 		// Invalid YAML.
 		{`"`, nil, true},
 		{`0ee1`, nil, true},
@@ -26,19 +28,6 @@ func TestNum_UnmarshalYAML(t *testing.T) {
 	}
 	for i, tt := range tests {
 		tt := tt
-		t.Run(fmt.Sprintf("Test%d", i+1), func(t *testing.T) {
-			a := require.New(t)
-
-			var val Num
-			err := yaml.Unmarshal([]byte(tt.data), &val)
-			if tt.wantErr {
-				a.Error(err)
-				t.Log("Input:", tt.data)
-				t.Log("Error:", err)
-				return
-			}
-			a.NoError(err)
-			a.Equal(tt.value, val)
-		})
+		t.Run(fmt.Sprintf("Test%d", i+1), testCustomEncodings(create, tt.data, tt.wantErr))
 	}
 }