Add support for oneOf inheritance

Bonsai.Sgen.Tests/DiscriminatorGenerationTests.cs

@@ -8,7 +8,7 @@ namespace Bonsai.Sgen.Tests
     public class DiscriminatorGenerationTests
     {
         [TestMethod]
-        public async Task GenerateDiscriminatorSchema_SerializerAnnotationsDeclareKnownTypes()
+        public async Task GenerateFromAnyOfDiscriminatorSchema_SerializerAnnotationsDeclareKnownTypes()
         {
             var schema = await JsonSchema.FromJsonAsync(@"
 {
@@ -48,7 +48,12 @@ public async Task GenerateDiscriminatorSchema_SerializerAnnotationsDeclareKnownT
       },
       ""Animal"": {
         ""type"": ""object"",
-        ""discriminator"": ""discriminator"",
+        ""discriminator"": {
+          ""propertyName"": ""discriminator"",
+          ""mapping"": {
+              ""DogType"": ""#/definitions/Dog""
+          }
+        },
         ""x-abstract"": true,
         ""additionalProperties"": false,
         ""required"": [
@@ -71,9 +76,109 @@ public async Task GenerateDiscriminatorSchema_SerializerAnnotationsDeclareKnownT
 ");
             var generator = TestHelper.CreateGenerator(schema);
             var code = generator.GenerateFile();
-            Assert.IsTrue(code.Contains("[JsonInheritanceAttribute(\"Dog\", typeof(Dog))]"));
+            Assert.IsTrue(code.Contains("[JsonInheritanceAttribute(\"DogType\", typeof(Dog))]"));
             Assert.IsTrue(code.Contains("[YamlDiscriminator(\"discriminator\")]"));
             CompilerTestHelper.CompileFromSource(code);
         }
+
+        [TestMethod]
+        public async Task GenerateFromOneOfDiscriminatorSchema_SerializerAnnotationsDeclareKnownTypes()
+        {
+            var schema = await JsonSchema.FromJsonAsync(@"
+{
+    ""$schema"": ""http://json-schema.org/draft-04/schema#"",
+    ""type"": ""object"",
+    ""title"": ""Container"",
+    ""properties"": {
+      ""Animals"": {
+        ""type"": ""array"",
+        ""items"": { ""$ref"": ""#/definitions/AnimalTypes"" }
+      }
+    },
+    ""definitions"": {
+      ""Dog"": {
+        ""type"": ""object"",
+        ""properties"": {
+          ""kind"": {
+            ""enum"": [ ""Dog"" ]
+          },
+          ""Bar"": {
+            ""type"": [
+              ""null"",
+              ""string""
+            ]
+          }
+        },
+        ""required"": [ ""kind"", ""Bar"" ],
+        ""allOf"": [
+          {
+            ""$ref"": ""#/definitions/Animal""
+          }
+        ]
+      },
+      ""Cat"": {
+        ""type"": ""object"",
+        ""properties"": {
+          ""kind"": {
+            ""enum"": [ ""Cat"" ]
+          },
+          ""Baz"": {
+            ""type"": [
+              ""null"",
+              ""string""
+            ]
+          }
+        },
+        ""required"": [ ""kind"", ""Baz"" ],
+        ""allOf"": [
+          {
+            ""$ref"": ""#/definitions/Animal""
+          }
+        ]
+      },
+      ""Animal"": {
+        ""type"": ""object"",
+        ""discriminator"": ""kind"",
+        ""x-abstract"": true,
+        ""required"": [
+          ""kind""
+        ],
+        ""properties"": {
+          ""Foo"": {
+            ""type"": [
+              ""null"",
+              ""string""
+            ]
+          },
+          ""kind"": {
+            ""type"": ""string""
+          }
+        }
+      },
+      ""AnimalTypes"": {
+          ""oneOf"": [
+            {
+              ""$ref"": ""#/definitions/Dog""
+            },
+            {
+              ""$ref"": ""#/definitions/Cat""
+            },
+            {
+              ""type"": ""null""
+            }
+          ]
+      }
+    }
+  }
+");
+            var generator = TestHelper.CreateGenerator(schema);
+            var code = generator.GenerateFile();
+            Assert.IsTrue(code.Contains("class Dog : Animal"), "Derived types do not inherit from base type.");
+            Assert.IsTrue(!code.Contains("public enum DogKind"), "Discriminator property is repeated in derived types.");
+            Assert.IsTrue(code.Contains("List<Animal> Animals"), "Container array element type does not match base type.");
+            Assert.IsTrue(code.Contains("[JsonInheritanceAttribute(\"Dog\", typeof(Dog))]"));
+            Assert.IsTrue(code.Contains("[YamlDiscriminator(\"kind\")]"));
+            CompilerTestHelper.CompileFromSource(code);
+        }
     }
 }

Bonsai.Sgen.Tests/TestHelper.cs

@@ -8,6 +8,7 @@ public static CSharpCodeDomGenerator CreateGenerator(
             JsonSchema schema,
             SerializerLibraries serializerLibraries = SerializerLibraries.YamlDotNet | SerializerLibraries.NewtonsoftJson)
         {
+            schema = schema.WithUniqueDiscriminatorProperties();
             var settings = new CSharpCodeDomGeneratorSettings
             {
                 Namespace = nameof(TestHelper),

Bonsai.Sgen/CSharpTypeResolver.cs

@@ -0,0 +1,93 @@
+using NJsonSchema;
+using NJsonSchema.CodeGeneration.CSharp;
+
+namespace Bonsai.Sgen
+{
+    internal class CSharpTypeResolver : NJsonSchema.CodeGeneration.CSharp.CSharpTypeResolver
+    {
+        readonly Dictionary<JsonSchema, JsonSchema> _baseTypeCache = new();
+
+        public CSharpTypeResolver(CSharpGeneratorSettings settings)
+            : base(settings)
+        {
+        }
+
+        public CSharpTypeResolver(CSharpGeneratorSettings settings, JsonSchema exceptionSchema)
+            : base(settings, exceptionSchema)
+        {
+        }
+
+        public override JsonSchema RemoveNullability(JsonSchema schema)
+        {
+            JsonSchema? selectedSchema = null;
+            foreach (JsonSchema o in schema.ActualSchema.OneOf)
+            {
+                if (o.IsNullable(SchemaType.JsonSchema))
+                {
+                    continue;
+                }
+
+                if (selectedSchema == null)
+                {
+                    selectedSchema = o;
+                }
+                else
+                {
+                    return ResolveBaseTypeSchema(schema.ActualSchema);
+                }
+            }
+
+            return selectedSchema ?? schema;
+        }
+
+        private JsonSchema ResolveBaseTypeSchema(JsonSchema schema)
+        {
+            if (!_baseTypeCache.TryGetValue(schema, out JsonSchema? baseSchema))
+            {
+                foreach (JsonSchema o in schema.OneOf)
+                {
+                    if (o.IsNullable(SchemaType.JsonSchema))
+                    {
+                        continue;
+                    }
+
+                    if (baseSchema == null)
+                    {
+                        baseSchema = o;
+                    }
+                    else
+                    {
+                        baseSchema = FindBestBaseSchema(baseSchema, o);
+                        if (baseSchema == null) break;
+                    }
+                }
+
+                baseSchema ??= JsonSchema.CreateAnySchema();
+                _baseTypeCache[schema] = baseSchema;
+            }
+
+            return baseSchema;
+        }
+
+        private static JsonSchema? FindBestBaseSchema(JsonSchema baseSchema, JsonSchema schema)
+        {
+            while (!IsAssignableFrom(baseSchema.ActualSchema, schema))
+            {
+                baseSchema = baseSchema.ActualSchema.InheritedSchema;
+                if (baseSchema == null) break;
+            }
+
+            return baseSchema;
+        }
+
+        private static bool IsAssignableFrom(JsonSchema schema, JsonSchema? typeSchema)
+        {
+            while (typeSchema?.ActualSchema != null && schema != typeSchema.ActualSchema)
+            {
+                typeSchema = typeSchema.ActualSchema.InheritedSchema;
+            }
+
+            return typeSchema?.ActualSchema != null;
+        }
+    }
+}

Bonsai.Sgen/JsonSchemaExtensions.cs

@@ -0,0 +1,45 @@
+using NJsonSchema;
+using NJsonSchema.CodeGeneration;
+using NJsonSchema.Visitors;
+
+namespace Bonsai.Sgen
+{
+    public static class JsonSchemaExtensions
+    {
+        public static JsonSchema WithUniqueDiscriminatorProperties(this JsonSchema schema)
+        {
+            var visitor = new DiscriminatorSchemaVisitor(schema);
+            visitor.Visit(schema);
+            return schema;
+        }
+
+        class DiscriminatorSchemaVisitor : JsonSchemaVisitorBase
+        {
+            public DiscriminatorSchemaVisitor(object rootObject)
+            {
+                RootObject = rootObject;
+            }
+
+            public object RootObject { get; }
+
+            protected override JsonSchema VisitSchema(JsonSchema schema, string path, string typeNameHint)
+            {
+                if (schema.DiscriminatorObject != null)
+                {
+                    foreach (var derivedSchema in schema.GetDerivedSchemas(RootObject).Keys)
+                    {
+                        foreach (var property in derivedSchema.Properties.Keys.ToList())
+                        {
+                            if (property == schema.Discriminator)
+                            {
+                                derivedSchema.Properties.Remove(property);
+                            }
+                        }
+                    }
+                }
+
+                return schema;
+            }
+        }
+    }
+}

Bonsai.Sgen/Program.cs

@@ -63,6 +63,7 @@ static async Task Main(string[] args)
                     SerializerLibraries = serializerLibraries
                 };
 
+                schema = schema.WithUniqueDiscriminatorProperties();
                 var generator = new CSharpCodeDomGenerator(schema, settings);
                 var code = generator.GenerateFile(generatorTypeName);
                 if (string.IsNullOrEmpty(outputFilePath))