Support inline oneOf discriminated properties

bonsai-rx · Jan 18, 2024 · 5fa58d2 · 5fa58d2
1 parent 68494ac
commit 5fa58d2
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/Bonsai.Sgen/JsonSchemaExtensions.cs b/Bonsai.Sgen/JsonSchemaExtensions.cs
@@ -15,17 +15,41 @@ public static JsonSchema WithUniqueDiscriminatorProperties(this JsonSchema schem
 
         class DiscriminatorSchemaVisitor : JsonSchemaVisitorBase
         {
-            public DiscriminatorSchemaVisitor(object rootObject)
+            public DiscriminatorSchemaVisitor(JsonSchema rootObject)
             {
                 RootObject = rootObject;
             }
 
-            public object RootObject { get; }
+            public JsonSchema RootObject { get; }
 
             protected override JsonSchema VisitSchema(JsonSchema schema, string path, string typeNameHint)
             {
                 if (schema.DiscriminatorObject != null)
                 {
+                    if (schema is JsonSchemaProperty schemaProperty)
+                    {
+                        if (!RootObject.Definitions.ContainsKey(typeNameHint))
+                        {
+                            var discriminatorSchema = new JsonSchema();
+                            discriminatorSchema.DiscriminatorObject = schemaProperty.DiscriminatorObject;
+                            discriminatorSchema.IsAbstract = schemaProperty.IsAbstract;
+                            foreach (var derivedSchema in schemaProperty.OneOf)
+                            {
+                                if (derivedSchema.IsNullable(SchemaType.JsonSchema))
+                                {
+                                    continue;
+                                }
+
+                                derivedSchema.ActualSchema.AllOf.Add(new JsonSchema { Reference = discriminatorSchema });
+                            }
+                            RootObject.Definitions.Add(typeNameHint, discriminatorSchema);
+                        }
+
+                        schemaProperty.DiscriminatorObject = null;
+                        schemaProperty.IsAbstract = false;
+                        return schema;
+                    }
+
                     foreach (var derivedSchema in schema.GetDerivedSchemas(RootObject).Keys)
                     {
                         foreach (var property in derivedSchema.Properties.Keys.ToList())