Merge pull request #419 from danielgtaylor/feat-discriminator

feat: basic schema discriminator support

schema.go

@@ -45,6 +45,28 @@ func deref(t reflect.Type) reflect.Type {
 	return t
 }
 
+// Discriminator object when request bodies or response payloads may be one of a
+// number of different schemas, can be used to aid in serialization,
+// deserialization, and validation. The discriminator is a specific object in a
+// schema which is used to inform the consumer of the document of an alternative
+// schema based on the value associated with it.
+type Discriminator struct {
+	// PropertyName in the payload that will hold the discriminator value.
+	// REQUIRED.
+	PropertyName string `yaml:"propertyName"`
+
+	// Mapping object to hold mappings between payload values and schema names or
+	// references.
+	Mapping map[string]string `yaml:"mapping,omitempty"`
+}
+
+func (d *Discriminator) MarshalJSON() ([]byte, error) {
+	return marshalJSON([]jsonFieldInfo{
+		{"propertyName", d.PropertyName, omitNever},
+		{"mapping", d.Mapping, omitEmpty},
+	}, nil)
+}
+
 // Schema represents a JSON Schema compatible with OpenAPI 3.1. It is extensible
 // with your own custom properties. It supports a subset of the full JSON Schema
 // spec, designed specifically for use with Go structs and to enable fast zero
@@ -99,6 +121,9 @@ type Schema struct {
 	AllOf []*Schema `yaml:"allOf,omitempty"`
 	Not   *Schema   `yaml:"not,omitempty"`
 
+	// OpenAPI specific fields
+	Discriminator *Discriminator `yaml:"discriminator,omitempty"`
+
 	patternRe     *regexp.Regexp  `yaml:"-"`
 	requiredMap   map[string]bool `yaml:"-"`
 	propertyNames []string        `yaml:"-"`
@@ -165,6 +190,7 @@ func (s *Schema) MarshalJSON() ([]byte, error) {
 		{"anyOf", s.AnyOf, omitEmpty},
 		{"allOf", s.AllOf, omitEmpty},
 		{"not", s.Not, omitEmpty},
+		{"discriminator", s.Discriminator, omitEmpty},
 	}, s.Extensions)
 }
 

schema_test.go

@@ -985,6 +985,55 @@ func TestCustomUnmarshalType(t *testing.T) {
 	assert.Equal(t, 0, o.Field.Value)
 }
 
+func TestMarshalDiscriminator(t *testing.T) {
+	s := &huma.Schema{
+		OneOf: []*huma.Schema{
+			{Type: "object", Properties: map[string]*huma.Schema{
+				"type": {Type: "string", Enum: []any{"foo"}},
+				"foo":  {Type: "string"},
+			}},
+			{Type: "object", Properties: map[string]*huma.Schema{
+				"type": {Type: "string", Enum: []any{"bar"}},
+				"bar":  {Type: "string"},
+			}},
+		},
+		Discriminator: &huma.Discriminator{
+			PropertyName: "type",
+			Mapping: map[string]string{
+				"foo": "#/components/schemas/Foo",
+				"bar": "#/components/schemas/Bar",
+			},
+		},
+	}
+
+	b, _ := json.Marshal(s)
+	assert.JSONEq(t, `{
+		"oneOf": [
+			{
+				"type": "object",
+				"properties": {
+					"type": {"type": "string", "enum": ["foo"]},
+					"foo": {"type": "string"}
+				}
+			},
+			{
+				"type": "object",
+				"properties": {
+					"type": {"type": "string", "enum": ["bar"]},
+					"bar": {"type": "string"}
+				}
+			}
+		],
+		"discriminator": {
+			"propertyName": "type",
+			"mapping": {
+				"foo": "#/components/schemas/Foo",
+				"bar": "#/components/schemas/Bar"
+			}
+		}
+	}`, string(b))
+}
+
 type BenchSub struct {
 	Visible bool      `json:"visible" default:"true"`
 	Metrics []float64 `json:"metrics" maxItems:"31"`