From 221c18e60a55b9847d158769f2735befb5a13db2 Mon Sep 17 00:00:00 2001
From: Simon Ser <contact@emersion.fr>
Date: Fri, 24 Feb 2023 16:53:44 +0100
Subject: [PATCH] Special-case nullable schema

---
 cmd/jsonschemagen/main.go | 35 +++++++++++++++++++++++++++++++++++
 1 file changed, 35 insertions(+)

diff --git a/cmd/jsonschemagen/main.go b/cmd/jsonschemagen/main.go
index dde8ebe..fdc01cd 100644
--- a/cmd/jsonschemagen/main.go
+++ b/cmd/jsonschemagen/main.go
@@ -128,6 +128,37 @@ func noAdditionalProps(schema *jsonschema.Schema) bool {
 	return schema.AdditionalProperties != nil && schema.AdditionalProperties.IsFalse()
 }
 
+// unwrapNullableSchema unwraps a schema in the form:
+//
+//	{
+//		"oneOf": {
+//			{ "type": "null" },
+//			<sub-schema>
+//		}
+//	}
+func unwrapNullableSchema(schema *jsonschema.Schema) (*jsonschema.Schema, bool) {
+	for _, choices := range [][]jsonschema.Schema{schema.AnyOf, schema.OneOf} {
+		if len(choices) != 2 {
+			continue
+		}
+
+		nullIndex := -1
+		for i, choice := range choices {
+			if len(choice.Type) == 1 && choice.Type[0] == jsonschema.TypeNull {
+				nullIndex = i
+				break
+			}
+		}
+		if nullIndex < 0 {
+			continue
+		}
+
+		otherIndex := (nullIndex + 1) % 2
+		return &choices[otherIndex], true
+	}
+	return nil, false
+}
+
 func generateSchemaType(schema *jsonschema.Schema, root *jsonschema.Schema, required bool) jen.Code {
 	if schema == nil {
 		return jen.Interface()
@@ -143,6 +174,10 @@ func generateSchemaType(schema *jsonschema.Schema, root *jsonschema.Schema, requ
 		return t
 	}
 
+	if subschema, ok := unwrapNullableSchema(schema); ok {
+		return jen.Op("*").Add(generateSchemaType(subschema, root, true))
+	}
+
 	switch schemaType(schema) {
 	case jsonschema.TypeNull:
 		return jen.Struct()