From 3b9f6e788bdff72d880c44e111850d31e383596b Mon Sep 17 00:00:00 2001
From: Louis Duchemin <ls.duchemin@gmail.com>
Date: Sat, 23 Mar 2024 09:56:13 +0100
Subject: [PATCH] Allow specifying type aliases to provide schema definitions

---
 registry.go      | 14 ++++++++++++++
 registry_test.go | 17 +++++++++++++++++
 2 files changed, 31 insertions(+)

diff --git a/registry.go b/registry.go
index 98f6cbab..ab828725 100644
--- a/registry.go
+++ b/registry.go
@@ -18,6 +18,7 @@ type Registry interface {
 	SchemaFromRef(ref string) *Schema
 	TypeFromRef(ref string) reflect.Type
 	Map() map[string]*Schema
+	RegisterTypeAlias(t reflect.Type, alias reflect.Type)
 }
 
 // DefaultSchemaNamer provides schema names for types. It uses the type name
@@ -62,10 +63,17 @@ type mapRegistry struct {
 	types   map[string]reflect.Type
 	seen    map[reflect.Type]bool
 	namer   func(reflect.Type, string) string
+	aliases map[reflect.Type]reflect.Type
 }
 
 func (r *mapRegistry) Schema(t reflect.Type, allowRef bool, hint string) *Schema {
 	t = deref(t)
+
+	alias, ok := r.aliases[t]
+	if ok {
+		return r.Schema(alias, allowRef, hint)
+	}
+
 	getsRef := t.Kind() == reflect.Struct
 	if t == timeType {
 		// Special case: time.Time is always a string.
@@ -134,6 +142,11 @@ func (r *mapRegistry) MarshalYAML() (interface{}, error) {
 	return r.schemas, nil
 }
 
+// RegisterTypeAlias(t, alias) makes the schema generator use the `alias` type instead of `t`.
+func (r *mapRegistry) RegisterTypeAlias(t reflect.Type, alias reflect.Type) {
+	r.aliases[t] = alias
+}
+
 // NewMapRegistry creates a new registry that stores schemas in a map and
 // returns references to them using the given prefix.
 func NewMapRegistry(prefix string, namer func(t reflect.Type, hint string) string) Registry {
@@ -142,6 +155,7 @@ func NewMapRegistry(prefix string, namer func(t reflect.Type, hint string) strin
 		schemas: map[string]*Schema{},
 		types:   map[string]reflect.Type{},
 		seen:    map[reflect.Type]bool{},
+		aliases: map[reflect.Type]reflect.Type{},
 		namer:   namer,
 	}
 }
diff --git a/registry_test.go b/registry_test.go
index 40c19f7b..6579c17f 100644
--- a/registry_test.go
+++ b/registry_test.go
@@ -53,3 +53,20 @@ func TestDefaultSchemaNamer(t *testing.T) {
 		})
 	}
 }
+
+func TestSchemaAlias(t *testing.T) {
+	type StringContainer struct {
+		Value string
+	}
+	type StructWithStringContainer struct {
+		Name StringContainer `json:"name,omitempty"`
+	}
+	type StructWithString struct {
+		Name string `json:"name,omitempty"`
+	}
+	registry := NewMapRegistry("#/components/schemas", DefaultSchemaNamer)
+	registry.RegisterTypeAlias(reflect.TypeOf(StringContainer{}), reflect.TypeOf(""))
+	schemaWithContainer := registry.Schema(reflect.TypeOf(StructWithStringContainer{}), false, "")
+	schemaWithString := registry.Schema(reflect.TypeOf(StructWithString{}), false, "")
+	assert.Equal(t, schemaWithString, schemaWithContainer)
+}