From 6feb227a4ccf5d60ddc720de292c0017a3d401c8 Mon Sep 17 00:00:00 2001
From: Rico Suter <mail@rsuter.com>
Date: Fri, 16 Mar 2018 17:57:40 +0100
Subject: [PATCH] Improved discriminator handling,
 https://github.com/RSuter/NSwag/issues/945

---
 .../Models/ClassTemplateModel.cs                 |  4 ++--
 .../TypeScriptTypeResolver.cs                    |  4 ++--
 .../Models/ClassTemplateModelBase.cs             | 16 +++++++++++-----
 src/NJsonSchema/JsonSchema4.cs                   | 15 +--------------
 4 files changed, 16 insertions(+), 23 deletions(-)

diff --git a/src/NJsonSchema.CodeGeneration.TypeScript/Models/ClassTemplateModel.cs b/src/NJsonSchema.CodeGeneration.TypeScript/Models/ClassTemplateModel.cs
index 31bd9924a..287ce4dab 100644
--- a/src/NJsonSchema.CodeGeneration.TypeScript/Models/ClassTemplateModel.cs
+++ b/src/NJsonSchema.CodeGeneration.TypeScript/Models/ClassTemplateModel.cs
@@ -49,10 +49,10 @@ public ClassTemplateModel(string typeName, string discriminatorName,
         public bool HasDiscriminator => !string.IsNullOrEmpty(_schema.Discriminator);
 
         /// <summary>Gets a value indicating whether the class or an inherited class has a discriminator property.</summary>
-        public bool HasBaseDiscriminator => !string.IsNullOrEmpty(_schema.BaseDiscriminator);
+        public bool HasBaseDiscriminator => _schema.BaseDiscriminator != null;
 
         /// <summary>Gets the class discriminator property name (may be defined in a inherited class).</summary>
-        public string BaseDiscriminator => _schema.BaseDiscriminator;
+        public string BaseDiscriminator => _schema.BaseDiscriminator?.PropertyName;
 
         /// <summary>Gets a value indicating whether the class has description.</summary>
         public bool HasDescription => !(_schema is JsonProperty) && !string.IsNullOrEmpty(_schema.Description);
diff --git a/src/NJsonSchema.CodeGeneration.TypeScript/TypeScriptTypeResolver.cs b/src/NJsonSchema.CodeGeneration.TypeScript/TypeScriptTypeResolver.cs
index bf985ab30..be3d39708 100644
--- a/src/NJsonSchema.CodeGeneration.TypeScript/TypeScriptTypeResolver.cs
+++ b/src/NJsonSchema.CodeGeneration.TypeScript/TypeScriptTypeResolver.cs
@@ -55,7 +55,7 @@ public override string Resolve(JsonSchema4 schema, bool isNullable, string typeN
         /// <returns>The result.</returns>
         public bool SupportsConstructorConversion(JsonSchema4 schema)
         {
-            return string.IsNullOrEmpty(schema?.ActualSchema.BaseDiscriminator);
+            return schema?.ActualSchema.BaseDiscriminator == null;
         }
 
         private string Resolve(JsonSchema4 schema, string typeNameHint, bool addInterfacePrefix)
@@ -124,7 +124,7 @@ private string ResolveString(JsonSchema4 schema, string typeNameHint)
                 if (schema.Format == JsonFormatStrings.TimeSpan)
                     return "string";
             }
-            else if (Settings.DateTimeType == TypeScriptDateTimeType.MomentJS || 
+            else if (Settings.DateTimeType == TypeScriptDateTimeType.MomentJS ||
                      Settings.DateTimeType == TypeScriptDateTimeType.OffsetMomentJS)
             {
                 if (schema.Format == JsonFormatStrings.Date)
diff --git a/src/NJsonSchema.CodeGeneration/Models/ClassTemplateModelBase.cs b/src/NJsonSchema.CodeGeneration/Models/ClassTemplateModelBase.cs
index 34c4df68a..484c31c8f 100644
--- a/src/NJsonSchema.CodeGeneration/Models/ClassTemplateModelBase.cs
+++ b/src/NJsonSchema.CodeGeneration/Models/ClassTemplateModelBase.cs
@@ -38,17 +38,23 @@ protected ClassTemplateModelBase(TypeResolverBase resolver, JsonSchema4 schema,
         /// <summary>Gets the derived class names (discriminator key/type name).</summary>
         public ICollection<DerivedClassModel> DerivedClasses => _schema
             .GetDerivedSchemas(_rootObject)
-            .Select(p => new DerivedClassModel(p.Value, _resolver.GetOrGenerateTypeName(p.Key, p.Value), p.Key.ActualTypeSchema.IsAbstract))
+            .Select(p => new DerivedClassModel(p.Value, p.Key, _schema.ActualSchema.BaseDiscriminator, _resolver))
             .ToList();
 
         /// <summary>The model of a derived class.</summary>
         public class DerivedClassModel
         {
-            internal DerivedClassModel(string discriminator, string className, bool isAbstract)
+            internal DerivedClassModel(string typeName, JsonSchema4 schema, OpenApiDiscriminator discriminator, TypeResolverBase resolver)
             {
-                ClassName = className;
-                IsAbstract = isAbstract;
-                Discriminator = !string.IsNullOrEmpty(discriminator) ? discriminator : className;
+                var mapping = discriminator.Mapping.SingleOrDefault(m => m.Value.ActualSchema == schema.ActualSchema);
+
+                ClassName = resolver.GetOrGenerateTypeName(schema, typeName);
+                IsAbstract = schema.ActualTypeSchema.IsAbstract;
+
+                Discriminator =
+                    mapping.Value != null ? mapping.Key :
+                    !string.IsNullOrEmpty(typeName) ? typeName :
+                    ClassName;
             }
 
             /// <summary>Gets the discriminator.</summary>
diff --git a/src/NJsonSchema/JsonSchema4.cs b/src/NJsonSchema/JsonSchema4.cs
index 853f93959..39db994be 100644
--- a/src/NJsonSchema/JsonSchema4.cs
+++ b/src/NJsonSchema/JsonSchema4.cs
@@ -279,20 +279,7 @@ public bool Inherits(JsonSchema4 schema)
 
         /// <summary>Gets the discriminator or discriminator of an inherited schema (or null).</summary>
         [JsonIgnore]
-        public string BaseDiscriminator
-        {
-            get
-            {
-                if (!string.IsNullOrEmpty(Discriminator))
-                    return Discriminator;
-
-                var baseDiscriminator = InheritedSchema?.ActualSchema.BaseDiscriminator;
-                if (!string.IsNullOrEmpty(baseDiscriminator))
-                    return baseDiscriminator;
-
-                return null;
-            }
-        }
+        public OpenApiDiscriminator BaseDiscriminator => DiscriminatorObject ?? InheritedSchema?.ActualSchema.BaseDiscriminator;
 
         /// <summary>Gets all properties of this schema (i.e. all direct properties and properties from the schemas in allOf which do not have a type).</summary>
         /// <remarks>Used for code generation.</remarks>