From a63c9cb8b23534ee92651e9824528b9c1c321cb2 Mon Sep 17 00:00:00 2001
From: "Daniel G. Taylor" <danielgtaylor@gmail.com>
Date: Mon, 19 Aug 2024 08:34:31 -0700
Subject: [PATCH] feat: treat encoding.TextUnmarshaler as string in schema

---
 registry.go    |  6 ++++++
 schema.go      | 11 ++++++++++-
 schema_test.go | 29 +++++++++++++++++++++++++++++
 3 files changed, 45 insertions(+), 1 deletion(-)

diff --git a/registry.go b/registry.go
index 222927c5..6ee69153 100644
--- a/registry.go
+++ b/registry.go
@@ -1,6 +1,7 @@
 package huma
 
 import (
+	"encoding"
 	"encoding/json"
 	"fmt"
 	"reflect"
@@ -92,6 +93,11 @@ func (r *mapRegistry) Schema(t reflect.Type, allowRef bool, hint string) *Schema
 		// Special case: type provides its own schema
 		getsRef = false
 	}
+	if _, ok := v.(encoding.TextUnmarshaler); ok {
+		// Special case: type can be unmarshalled from text so will be a `string`
+		// and doesn't need a ref. This simplifies the schema a little bit.
+		getsRef = false
+	}
 
 	name := r.namer(origType, hint)
 
diff --git a/schema.go b/schema.go
index 98d208b9..847b2283 100644
--- a/schema.go
+++ b/schema.go
@@ -1,6 +1,7 @@
 package huma
 
 import (
+	"encoding"
 	"encoding/json"
 	"errors"
 	"fmt"
@@ -686,7 +687,7 @@ func schemaFromType(r Registry, t reflect.Type) *Schema {
 		return custom
 	}
 
-	// Handle special cases.
+	// Handle special cases for known stdlib types.
 	switch t {
 	case timeType:
 		return &Schema{Type: TypeString, Nullable: isPointer, Format: "date-time"}
@@ -700,6 +701,14 @@ func schemaFromType(r Registry, t reflect.Type) *Schema {
 		return &Schema{}
 	}
 
+	if _, ok := v.(encoding.TextUnmarshaler); ok {
+		// Special case: types that implement encoding.TextUnmarshaler are able to
+		// be loaded from plain text, and so should be treated as strings.
+		// This behavior can be overidden by implementing `huma.SchemaProvider`
+		// and returning a custom schema.
+		return &Schema{Type: TypeString, Nullable: isPointer}
+	}
+
 	minZero := 0.0
 	switch t.Kind() {
 	case reflect.Bool:
diff --git a/schema_test.go b/schema_test.go
index 6f8a44ff..8f880d6b 100644
--- a/schema_test.go
+++ b/schema_test.go
@@ -2,6 +2,7 @@ package huma_test
 
 import (
 	"bytes"
+	"encoding"
 	"encoding/json"
 	"math/bits"
 	"net"
@@ -1172,6 +1173,34 @@ func TestSchemaGenericNamingFromModule(t *testing.T) {
 	}`, string(b))
 }
 
+type MyDate time.Time
+
+func (d *MyDate) UnmarshalText(data []byte) error {
+	t, err := time.Parse(time.RFC3339, string(data))
+	if err != nil {
+		return err
+	}
+	*d = MyDate(t)
+	return nil
+}
+
+var _ encoding.TextUnmarshaler = (*MyDate)(nil)
+
+func TestCustomDateType(t *testing.T) {
+	type O struct {
+		Date MyDate `json:"date"`
+	}
+
+	var o O
+	err := json.Unmarshal([]byte(`{"date": "2022-01-01T00:00:00Z"}`), &o)
+	require.NoError(t, err)
+	assert.Equal(t, MyDate(time.Date(2022, 1, 1, 0, 0, 0, 0, time.UTC)), o.Date)
+
+	r := huma.NewMapRegistry("#/components/schemas/", huma.DefaultSchemaNamer)
+	s := r.Schema(reflect.TypeOf(o), false, "")
+	assert.Equal(t, "string", s.Properties["date"].Type)
+}
+
 type OmittableNullable[T any] struct {
 	Sent  bool
 	Null  bool