Merge pull request #292 from csdahlberg/master

Added support for KnownType attributes that specify the name of a method
RicoSuter · Jan 9, 2017 · 380a1f4 · 380a1f4
2 parents f99fa8f + 5c4478b
commit 380a1f4
Show file tree

Hide file tree

Showing 2 changed files with 61 additions and 4 deletions.
diff --git a/src/NJsonSchema.Tests/Generation/KnownTypeGenerationTests.cs b/src/NJsonSchema.Tests/Generation/KnownTypeGenerationTests.cs
@@ -20,9 +20,27 @@ public class Person
         {
         }
 
+        public class Pen
+        {
+        }
+
+        public class Pencil
+        {
+        }
+
+        [KnownType("GetKnownTypes")]
+        public class WritingInstrument
+        {
+            public static Type[] GetKnownTypes()
+            {
+                return new[] { typeof(Pen), typeof(Pencil) };
+            }
+        }
+
         public class Container
         {
             public Person Person { get; set; }
+            public WritingInstrument WritingInstrument { get; set; }
         }
 
         [TestMethod]
@@ -37,5 +55,23 @@ public async Task When_KnownType_attribute_exists_then_specified_classes_are_als
             //// Assert
             Assert.IsTrue(schema.Definitions.Any(s => s.Key == "Teacher"));
         }
+
+        public async Task ReproAsync()
+        {
+            var schema = await JsonSchema4.FromTypeAsync<Container>();
+        }
+        [TestMethod]
+        public async Task When_KnownType_attribute_includes_method_name_then_specified_classes_are_also_generated()
+        {
+            //// Arrange
+
+            //// Act
+            var schema = await JsonSchema4.FromTypeAsync<Container>();
+            var schemaData = schema.ToJson();
+
+            //// Assert
+            Assert.IsTrue(schema.Definitions.Any(s => s.Key == "Pen"));
+            Assert.IsTrue(schema.Definitions.Any(s => s.Key == "Pencil"));
+        }
     }
 }
diff --git a/src/NJsonSchema/Generation/JsonSchemaGenerator.cs b/src/NJsonSchema/Generation/JsonSchemaGenerator.cs
@@ -305,11 +305,32 @@ private async Task GenerateKnownTypesAsync(Type type, JsonSchemaResolver schemaR
         {
             foreach (dynamic knownTypeAttribute in type.GetTypeInfo().GetCustomAttributes().Where(a => a.GetType().Name == "KnownTypeAttribute"))
             {
-                var typeDescription = JsonObjectTypeDescription.FromType(knownTypeAttribute.Type, null, Settings.DefaultEnumHandling);
-                var isIntegerEnum = typeDescription.Type == JsonObjectType.Integer;
+                if (knownTypeAttribute.Type != null)
+                {
+                    var typeDescription = JsonObjectTypeDescription.FromType(knownTypeAttribute.Type, null, Settings.DefaultEnumHandling);
+                    var isIntegerEnum = typeDescription.Type == JsonObjectType.Integer;
+
+                    if (!schemaResolver.HasSchema(knownTypeAttribute.Type, isIntegerEnum))
+                        await GenerateAsync(knownTypeAttribute.Type, schemaResolver).ConfigureAwait(false);
+                }
+                else if (!string.IsNullOrWhiteSpace(knownTypeAttribute.MethodName))
+                {
+                    var methodInfo = type.GetRuntimeMethod((string)knownTypeAttribute.MethodName, new Type[0]);
+                    var knownTypes = methodInfo.Invoke(null, null) as Type[];
+
+                    foreach (var knownType in knownTypes)
+                    {
+                        var typeDescription = JsonObjectTypeDescription.FromType(knownType, null, Settings.DefaultEnumHandling);
+                        var isIntegerEnum = typeDescription.Type == JsonObjectType.Integer;
 
-                if (!schemaResolver.HasSchema(knownTypeAttribute.Type, isIntegerEnum))
-                    await GenerateAsync(knownTypeAttribute.Type, schemaResolver).ConfigureAwait(false);
+                        if (!schemaResolver.HasSchema(knownType, isIntegerEnum))
+                            await GenerateAsync(knownType, schemaResolver).ConfigureAwait(false);
+                    }
+                }
+                else
+                {
+                    throw new ArgumentException($"A KnownType attribute on {type.FullName} does not specify a type or a method name.", nameof(type));
+                }
             }
         }