Repository: ogen-go/ogen
Pull request #416 (2022-06-22 13:56:59)

feat(jsonschema): add simple tool for JSON schema inference


Patch:
diff --git a/jsonschema/infer.go b/jsonschema/infer.go
new file mode 100644
index 00000000..4ab5c898
--- /dev/null
+++ b/jsonschema/infer.go
@@ -0,0 +1,188 @@
+package jsonschema
+
+import (
+	"sort"
+
+	"github.com/go-faster/errors"
+	"github.com/go-faster/jx"
+)
+
+// Infer returns a JSON Schema that is inferred from the given JSON.
+type Infer struct {
+	target RawSchema
+}
+
+// Target returns the target schema.
+func (i Infer) Target() RawSchema {
+	return i.target
+}
+
+// Apply applies given data to the schema state.
+func (i *Infer) Apply(data []byte) error {
+	return apply(&i.target, jx.DecodeBytes(data))
+}
+
+func applyType(s *RawSchema, tt string) {
+	if hasType(s, tt) {
+		return
+	}
+	if len(s.OneOf) > 0 {
+		s.OneOf = append(s.OneOf, &RawSchema{Type: tt})
+		return
+	}
+
+	if s.Type == "" {
+		s.Type = tt
+		return
+	}
+
+	old := *s
+	*s = RawSchema{
+		OneOf: []*RawSchema{
+			&old,
+			{Type: tt},
+		},
+	}
+}
+
+func hasType(s *RawSchema, tt string) bool {
+	if s.Type == tt {
+		return true
+	}
+	for _, v := range s.OneOf {
+		if v.Type == tt {
+			return true
+		}
+	}
+	return false
+}
+
+func replaceType(s *RawSchema, from, to string) bool {
+	if s.Type == from {
+		s.Type = to
+		return true
+	}
+	for _, v := range s.OneOf {
+		if v.Type == from {
+			v.Type = to
+			return true
+		}
+	}
+	return false
+}
+
+func apply(s *RawSchema, d *jx.Decoder) error {
+	switch tt := d.Next(); tt {
+	case jx.String:
+		applyType(s, "string")
+		return d.Skip()
+	case jx.Number:
+		n, err := d.Num()
+		if err != nil {
+			return err
+		}
+		if n.IsInt() && !hasType(s, "number") {
+			applyType(s, "integer")
+			return nil
+		}
+		if replaceType(s, "integer", "number") {
+			return nil
+		}
+		applyType(s, "number")
+		return nil
+	case jx.Null:
+		s.Nullable = true
+		return d.Skip()
+	case jx.Bool:
+		applyType(s, "boolean")
+		return d.Skip()
+	case jx.Array:
+		applyType(s, "array")
+
+		i := 0
+		return d.Arr(func(d *jx.Decoder) error {
+			if s.Items == nil {
+				s.Items = new(RawSchema)
+			}
+			if err := apply(s.Items, d); err != nil {
+				return errors.Wrapf(err, "apply item %d", i)
+			}
+			i++
+			return nil
+		})
+	case jx.Object:
+		applyType(s, "object")
+
+		// Set s.Properties to non-nil slice to mark that it is not first apply.
+		firstApply := s.Properties == nil
+		if firstApply {
+			s.Properties = RawProperties{}
+		}
+
+		// Collect existing properties.
+		props := map[string]*RawSchema{}
+		for _, prop := range s.Properties {
+			props[prop.Name] = prop.Schema
+		}
+
+		// Collect required properties.
+		required := map[string]struct{}{}
+		for _, key := range s.Required {
+			required[key] = struct{}{}
+		}
+
+		this := map[string]struct{}{}
+		if err := d.Obj(func(d *jx.Decoder, key string) error {
+			this[key] = struct{}{}
+
+			if err := func() error {
+				if prop, ok := props[key]; ok {
+					return apply(prop, d)
+				}
+
+				// If it is the first apply, mark property as required.
+				if firstApply {
+					required[key] = struct{}{}
+				}
+
+				prop := new(RawSchema)
+				if err := apply(prop, d); err != nil {
+					return err
+				}
+				s.Properties = append(s.Properties, RawProperty{
+					Name:   key,
+					Schema: prop,
+				})
+				return nil
+			}(); err != nil {
+				return errors.Wrapf(err, "apply property %q", key)
+			}
+			return nil
+		}); err != nil {
+			return errors.Wrap(err, "collect properties")
+		}
+
+		// Delete required properties that are not in this object.
+		for key := range required {
+			if _, ok := this[key]; !ok {
+				delete(required, key)
+			}
+		}
+
+		// Write required properties.
+		s.Required = s.Required[:0]
+		for key := range required {
+			s.Required = append(s.Required, key)
+		}
+
+		// Sort fields to make output deterministic.
+		sort.Strings(s.Required)
+		sort.SliceStable(s.Properties, func(i, j int) bool {
+			return s.Properties[i].Name < s.Properties[j].Name
+		})
+
+		return nil
+	default:
+		return errors.Errorf("invalid type: %s", tt)
+	}
+}


Test patch:
diff --git a/jsonschema/infer_test.go b/jsonschema/infer_test.go
new file mode 100644
index 00000000..c8b3811f
--- /dev/null
+++ b/jsonschema/infer_test.go
@@ -0,0 +1,127 @@
+package jsonschema
+
+import (
+	"fmt"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestInfer_Apply(t *testing.T) {
+	t.Run("Bad", func(t *testing.T) {
+		for _, input := range []string{
+			``,
+			`r`,
+			`-..`,
+			`{`,
+			`[`,
+			`[{`,
+			`[{}`,
+			`{"foo": [-..]}`,
+		} {
+			var i Infer
+			require.Errorf(t, i.Apply([]byte(input)), "input:\n%s", input)
+		}
+	})
+
+	tests := []struct {
+		result RawSchema
+		inputs []string
+	}{
+		{RawSchema{Type: "integer"}, []string{"1", "2", "3"}},
+		{RawSchema{Type: "number"}, []string{"1", "2.0", "3"}},
+		{RawSchema{Type: "number"}, []string{"2.0"}},
+		{RawSchema{Type: "number", Nullable: true}, []string{"2.0", "null"}},
+
+		{RawSchema{Type: "boolean"}, []string{"true", "false"}},
+		{RawSchema{Type: "boolean", Nullable: true}, []string{"true", "null"}},
+
+		{RawSchema{Type: "array"}, []string{"[]"}},
+		{RawSchema{
+			Type:  "array",
+			Items: &RawSchema{Type: "integer"},
+		}, []string{"[1]"}},
+		{RawSchema{
+			Type:  "array",
+			Items: &RawSchema{Type: "number"},
+		}, []string{"[1, 10, 5, 0.5]"}},
+		{RawSchema{
+			Type: "array",
+			Items: &RawSchema{
+				OneOf: []*RawSchema{
+					{Type: "integer"},
+					{Type: "boolean"},
+					{Type: "string"},
+				},
+			},
+		}, []string{`[1, true, "foo"]`}},
+
+		{RawSchema{Type: "object", Properties: RawProperties{}}, []string{
+			`{}`,
+		}},
+		{RawSchema{
+			Type: "object",
+			Properties: RawProperties{
+				{"foo", &RawSchema{Type: "integer"}},
+			},
+		}, []string{
+			`{}`,
+			`{"foo": 1}`,
+			`{"foo": 2}`,
+			`{"foo": 3}`,
+		}},
+		{RawSchema{
+			Type:     "object",
+			Required: []string{"foo"},
+			Properties: RawProperties{
+				{"bar", &RawSchema{Type: "string"}},
+				{"foo", &RawSchema{Type: "integer"}},
+			},
+		}, []string{
+			`{"foo": 1}`,
+			`{"foo": 5}`,
+			`{"foo": 2, "bar": "baz"}`,
+		}},
+		{RawSchema{
+			Type:     "object",
+			Required: []string{"required", "required_nullable"},
+			Properties: RawProperties{
+				{"optional", &RawSchema{Type: "integer"}},
+				{"optional_nullable", &RawSchema{Type: "integer", Nullable: true}},
+				{"required", &RawSchema{Type: "integer"}},
+				{"required_nullable", &RawSchema{Type: "integer", Nullable: true}},
+			},
+		}, []string{
+			`{"required": 10, "required_nullable": null, "optional": 10, "optional_nullable": null}`,
+			`{"required": 10, "required_nullable": 10}`,
+			`{"required": 10, "required_nullable": 10, "optional_nullable": 10}`,
+		}},
+
+		{RawSchema{Nullable: true}, []string{"null"}},
+		{RawSchema{
+			OneOf: []*RawSchema{
+				{Type: "boolean"},
+				{Type: "string"},
+				{Type: "number"},
+			},
+		}, []string{"true", `"foo"`, "10", "1.0"}},
+		{RawSchema{
+			OneOf: []*RawSchema{
+				{Type: "boolean"},
+				{Type: "string"},
+				{Type: "number"},
+			},
+		}, []string{"true", `"foo"`, "1.0", "10"}},
+	}
+	for i, tt := range tests {
+		tt := tt
+		t.Run(fmt.Sprintf("Test%d", i+1), func(t *testing.T) {
+			a := require.New(t)
+			var i Infer
+			for _, input := range tt.inputs {
+				a.NoError(i.Apply([]byte(input)))
+			}
+			a.Equal(tt.result, i.Target())
+		})
+	}
+}