Merge pull request #36 from bonsai-rx/discriminator-dev

Add support for deep oneOf inheritance
bonsai-rx · Jan 24, 2024 · a11b1a8 · a11b1a8
2 parents 9ad689f + b493bb9
commit a11b1a8
Show file tree

Hide file tree

Showing 7 changed files with 85 additions and 36 deletions.
diff --git a/Bonsai.Sgen.Tests/DiscriminatorGenerationTests.cs b/Bonsai.Sgen.Tests/DiscriminatorGenerationTests.cs
@@ -1,5 +1,6 @@
 using Microsoft.VisualStudio.TestTools.UnitTesting;
 using NJsonSchema;
+using System.Collections.Generic;
 using System.Linq;
 using System.Threading.Tasks;
 
@@ -214,10 +215,49 @@ public void GenerateFromArrayItemDiscriminatorRef_EnsureFallbackDiscriminatorBas
             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> Animal"), "Container element type does not match base type.");
+            Assert.IsTrue(code.Contains("List<Animal> Animals"), "Container element type does not match base type.");
             Assert.IsTrue(code.Contains("[JsonInheritanceAttribute(\"Dog\", typeof(Dog))]"));
             AssertDiscriminatorAttribute(code, serializerLibraries, "kind");
             CompilerTestHelper.CompileFromSource(code);
         }
+
+        [TestMethod]
+        [DataRow(SerializerLibraries.YamlDotNet)]
+        [DataRow(SerializerLibraries.NewtonsoftJson)]
+        [DataRow(SerializerLibraries.NewtonsoftJson | SerializerLibraries.YamlDotNet)]
+        public void GenerateFromSubDiscriminatorSchemas_InheritanceHierarchyIsPreserved(SerializerLibraries serializerLibraries)
+        {
+            var subTypeSchemas = SchemaTestHelper.CreateDerivedSchemas("fur", "long", "short");
+            var subDiscriminator = SchemaTestHelper.CreateDiscriminatorSchema("fur", subTypeSchemas);
+            var derivedSchemas = SchemaTestHelper.CreateDerivedSchemas("type", "cat", "dog");
+            derivedSchemas[0].Value.DiscriminatorObject = subDiscriminator.DiscriminatorObject;
+            foreach (var subSchema in subDiscriminator.OneOf)
+            {
+                derivedSchemas[0].Value.OneOf.Add(subSchema);
+            }
+            var discriminator = SchemaTestHelper.CreateDiscriminatorSchema("type", derivedSchemas);
+            var schema = SchemaTestHelper.CreateContainerSchema(new Dictionary<string, JsonSchema>
+            {
+                { "LongFurCat", subTypeSchemas[0].Value },
+                { "Cat", derivedSchemas[0].Value },
+                { "Dog", derivedSchemas[1].Value },
+                { "Animal", discriminator },
+                { "ShortFurCat", subTypeSchemas[1].Value }
+            });
+            schema.Properties.Add("Animals", new()
+            {
+                Type = JsonObjectType.Array,
+                Item = new() { Reference = discriminator }
+            });
+
+            var generator = TestHelper.CreateGenerator(schema, serializerLibraries);
+            var code = generator.GenerateFile();
+            Assert.IsTrue(code.Contains("class Cat : Animal"), "Derived types do not inherit from base type.");
+            Assert.IsTrue(!code.Contains("public enum LongFurCatFur"), "Discriminator property is repeated in derived types.");
+            Assert.IsTrue(code.Contains("List<Animal> Animals"), "Container element type does not match base type.");
+            Assert.IsTrue(code.Contains("[JsonInheritanceAttribute(\"dog\", typeof(Dog))]"));
+            AssertDiscriminatorAttribute(code, serializerLibraries, "type");
+            CompilerTestHelper.CompileFromSource(code);
+        }
     }
 }
diff --git a/Bonsai.Sgen.Tests/TestHelper.cs b/Bonsai.Sgen.Tests/TestHelper.cs
@@ -8,7 +8,7 @@ public static CSharpCodeDomGenerator CreateGenerator(
             JsonSchema schema,
             SerializerLibraries serializerLibraries = SerializerLibraries.YamlDotNet | SerializerLibraries.NewtonsoftJson)
         {
-            schema = schema.WithUniqueDiscriminatorProperties();
+            schema = schema.WithResolvedDiscriminatorInheritance();
             var settings = new CSharpCodeDomGeneratorSettings
             {
                 Namespace = nameof(TestHelper),

diff --git a/Bonsai.Sgen/CSharpClassTemplate.cs b/Bonsai.Sgen/CSharpClassTemplate.cs
@@ -4,6 +4,7 @@
 using System.Text;
 using System.Xml.Serialization;
 using Newtonsoft.Json;
+using NJsonSchema;
 using NJsonSchema.Converters;
 using YamlDotNet.Serialization;
 
@@ -30,7 +31,7 @@ public override void BuildType(CodeTypeDeclaration type)
             var jsonSerializer = Settings.SerializerLibraries.HasFlag(SerializerLibraries.NewtonsoftJson);
             var yamlSerializer = Settings.SerializerLibraries.HasFlag(SerializerLibraries.YamlDotNet);
             if (Model.IsAbstract) type.TypeAttributes |= System.Reflection.TypeAttributes.Abstract;
-            if (Model.HasDiscriminator)
+            if (Model.Schema.DiscriminatorObject is OpenApiDiscriminator discriminator)
             {
                 if (jsonSerializer || yamlSerializer)
                 {
@@ -39,13 +40,13 @@ public override void BuildType(CodeTypeDeclaration type)
                         type.CustomAttributes.Add(new CodeAttributeDeclaration(
                             new CodeTypeReference(typeof(JsonConverter)),
                             new CodeAttributeArgument(new CodeTypeOfExpression(nameof(JsonInheritanceConverter))),
-                            new CodeAttributeArgument(new CodePrimitiveExpression(Model.Discriminator))));
+                            new CodeAttributeArgument(new CodePrimitiveExpression(discriminator.PropertyName))));
                     }
                     if (yamlSerializer)
                     {
                         type.CustomAttributes.Add(new CodeAttributeDeclaration(
                             new CodeTypeReference("YamlDiscriminator"),
-                            new CodeAttributeArgument(new CodePrimitiveExpression(Model.Discriminator))));
+                            new CodeAttributeArgument(new CodePrimitiveExpression(discriminator.PropertyName))));
                     }
 
                     foreach (var derivedModel in Model.DerivedClasses)

diff --git a/Bonsai.Sgen/CSharpCodeDomGenerator.cs b/Bonsai.Sgen/CSharpCodeDomGenerator.cs
@@ -105,7 +105,7 @@ public override IEnumerable<CodeArtifact> GenerateTypes()
                               let classType = type as CSharpClassCodeArtifact
                               where classType != null
                               select classType).ToList();
-            var discriminatorTypes = classTypes.Where(modelType => modelType.Model.HasDiscriminator).ToList();
+            var discriminatorTypes = classTypes.Where(modelType => modelType.Model.Schema.DiscriminatorObject != null).ToList();
             foreach (var type in discriminatorTypes)
             {
                 var matchTemplate = new CSharpTypeMatchTemplate(type, _provider, _options, Settings);

diff --git a/Bonsai.Sgen/CSharpTypeResolver.cs b/Bonsai.Sgen/CSharpTypeResolver.cs
@@ -12,11 +12,6 @@ public CSharpTypeResolver(CSharpGeneratorSettings settings)
         {
         }
 
-        public CSharpTypeResolver(CSharpGeneratorSettings settings, JsonSchema exceptionSchema)
-            : base(settings, exceptionSchema)
-        {
-        }
-
         public override JsonSchema RemoveNullability(JsonSchema schema)
         {
             JsonSchema? selectedSchema = null;

diff --git a/Bonsai.Sgen/JsonSchemaExtensions.cs b/Bonsai.Sgen/JsonSchemaExtensions.cs
@@ -1,21 +1,22 @@
 using NJsonSchema;
-using NJsonSchema.CodeGeneration;
 using NJsonSchema.Visitors;
 
 namespace Bonsai.Sgen
 {
     internal static class JsonSchemaExtensions
     {
-        public static JsonSchema WithUniqueDiscriminatorProperties(this JsonSchema schema)
+        public static JsonSchema WithResolvedDiscriminatorInheritance(this JsonSchema schema)
         {
-            var visitor = new DiscriminatorSchemaVisitor(schema);
-            visitor.Visit(schema);
+            var discriminatorVisitor = new DiscriminatorSchemaVisitor(schema);
+            var derivedDiscriminatorVisitor = new DerivedDiscriminatorSchemaVisitor();
+            discriminatorVisitor.Visit(schema);
+            derivedDiscriminatorVisitor.Visit(schema);
             return schema;
         }
 
         class DiscriminatorSchemaVisitor : JsonSchemaVisitorBase
         {
-            readonly Dictionary<JsonSchema, string> reverseTypeNameLookup = new();
+            readonly Dictionary<JsonSchema, string> definitionTypeNameLookup = new();
 
             public DiscriminatorSchemaVisitor(JsonSchema rootObject)
             {
@@ -34,7 +35,11 @@ private void ResolveOneOfInheritance(JsonSchema schema, JsonSchema baseSchema)
                         continue;
                     }
 
-                    derivedSchema.ActualSchema.AllOf.Add(new JsonSchema { Reference = baseSchema });
+                    var actualSchema = derivedSchema.ActualSchema;
+                    if (!actualSchema.AllOf.Any(schema => schema.Reference == baseSchema))
+                    {
+                        actualSchema.AllOf.Add(new JsonSchema { Reference = baseSchema });
+                    }
                 }
             }
 
@@ -43,15 +48,16 @@ protected override JsonSchema VisitSchema(JsonSchema schema, string path, string
                 var actualSchema = schema.ActualSchema;
                 if (actualSchema.DiscriminatorObject != null)
                 {
-                    if (schema is JsonSchemaProperty || schema.ParentSchema?.Item == schema)
+                    var isDefinition = definitionTypeNameLookup.TryGetValue(actualSchema, out _);
+                    if (schema is JsonSchemaProperty || schema.ParentSchema?.Item == schema || isDefinition)
                     {
-                        if (string.IsNullOrEmpty(typeNameHint) &&
-                            !reverseTypeNameLookup.TryGetValue(actualSchema, out typeNameHint))
+                        var discriminatorSchema = isDefinition ? actualSchema : null;
+                        if (string.IsNullOrEmpty(typeNameHint))
                         {
                             typeNameHint = "Anonymous";
                         }
 
-                        if (!RootObject.Definitions.TryGetValue(typeNameHint, out JsonSchema? discriminatorSchema))
+                        if (discriminatorSchema == null && !RootObject.Definitions.TryGetValue(typeNameHint, out discriminatorSchema))
                         {
                             discriminatorSchema = new JsonSchema();
                             discriminatorSchema.DiscriminatorObject = actualSchema.DiscriminatorObject;
@@ -64,23 +70,13 @@ protected override JsonSchema VisitSchema(JsonSchema schema, string path, string
                             if (discriminatorSchema.OneOf.Count > 0)
                             {
                                 ResolveOneOfInheritance(discriminatorSchema, discriminatorSchema);
-                                discriminatorSchema.OneOf.Clear();
                             }
                         }
 
-                        schema.DiscriminatorObject = null;
-                        schema.IsAbstract = false;
-                        return schema;
-                    }
-
-                    foreach (var derivedSchema in schema.GetDerivedSchemas(RootObject).Keys)
-                    {
-                        foreach (var property in derivedSchema.Properties.Keys.ToList())
+                        if (!isDefinition)
                         {
-                            if (property == schema.Discriminator)
-                            {
-                                derivedSchema.Properties.Remove(property);
-                            }
+                            actualSchema.DiscriminatorObject = null;
+                            actualSchema.IsAbstract = false;
                         }
                     }
                 }
@@ -111,7 +107,7 @@ private void VisitDefinitions(JsonSchema schema)
                 {
                     foreach (var definition in schema.Definitions)
                     {
-                        reverseTypeNameLookup[definition.Value] = definition.Key;
+                        definitionTypeNameLookup[definition.Value] = definition.Key;
                         VisitDefinitions(definition.Value);
                     }
                 }
@@ -139,5 +135,22 @@ private void VisitDefinitions(IDictionary<string, JsonSchemaProperty> dictionary
                 }
             }
         }
+
+        class DerivedDiscriminatorSchemaVisitor : JsonSchemaVisitorBase
+        {
+            protected override JsonSchema VisitSchema(JsonSchema schema, string path, string typeNameHint)
+            {
+                foreach (var baseSchema in schema.AllInheritedSchemas)
+                {
+                    var discriminatorSchema = baseSchema.DiscriminatorObject;
+                    if (discriminatorSchema != null)
+                    {
+                        schema.Properties.Remove(discriminatorSchema.PropertyName);
+                    }
+                }
+
+                return schema;
+            }
+        }
     }
 }
diff --git a/Bonsai.Sgen/Program.cs b/Bonsai.Sgen/Program.cs
@@ -63,7 +63,7 @@ static async Task Main(string[] args)
                     SerializerLibraries = serializerLibraries
                 };
 
-                schema = schema.WithUniqueDiscriminatorProperties();
+                schema = schema.WithResolvedDiscriminatorInheritance();
                 var generator = new CSharpCodeDomGenerator(schema, settings);
                 var code = generator.GenerateFile(generatorTypeName);
                 if (string.IsNullOrEmpty(outputFilePath))