Merge pull request #368 from ionite34/enum-converter-fixes

StabilityMatrix.Core/Converters/Json/DefaultUnknownEnumConverter.cs

@@ -1,38 +1,95 @@
-using System.Text.Json;
+using System.Diagnostics.CodeAnalysis;
+using System.Reflection;
+using System.Runtime.Serialization;
+using System.Text.Json;
 using System.Text.Json.Serialization;
-using StabilityMatrix.Core.Extensions;
 
 namespace StabilityMatrix.Core.Converters.Json;
 
-public class DefaultUnknownEnumConverter<T> : JsonConverter<T>
+public class DefaultUnknownEnumConverter<
+    [DynamicallyAccessedMembers(DynamicallyAccessedMemberTypes.PublicFields)] T
+> : JsonConverter<T>
     where T : Enum
 {
+    /// <summary>
+    /// Lazy initialization for <see cref="EnumMemberValues"/>.
+    /// </summary>
+    private readonly Lazy<Dictionary<string, T>> _enumMemberValuesLazy =
+        new(
+            () =>
+                typeof(T)
+                    .GetFields()
+                    .Where(field => field.IsStatic)
+                    .Select(
+                        field =>
+                            new
+                            {
+                                FieldName = field.Name,
+                                FieldValue = (T)field.GetValue(null)!,
+                                EnumMemberValue = field
+                                    .GetCustomAttributes<EnumMemberAttribute>(false)
+                                    .FirstOrDefault()
+                                    ?.Value?.ToString()
+                            }
+                    )
+                    .ToDictionary(x => x.EnumMemberValue ?? x.FieldName, x => x.FieldValue)
+        );
+
+    /// <summary>
+    /// Gets a dictionary of enum member values, keyed by the EnumMember attribute value, or the field name if no EnumMember attribute is present.
+    /// </summary>
+    private Dictionary<string, T> EnumMemberValues => _enumMemberValuesLazy.Value;
+
+    /// <summary>
+    /// Lazy initialization for <see cref="EnumMemberNames"/>.
+    /// </summary>
+    private readonly Lazy<Dictionary<T, string>> _enumMemberNamesLazy;
+
+    /// <summary>
+    /// Gets a dictionary of enum member names, keyed by the enum member value.
+    /// </summary>
+    private Dictionary<T, string> EnumMemberNames => _enumMemberNamesLazy.Value;
+
+    /// <summary>
+    /// Gets the value of the "Unknown" enum member, or the 0 value if no "Unknown" member is present.
+    /// </summary>
+    private T UnknownValue =>
+        EnumMemberValues.TryGetValue("Unknown", out var res) ? res : (T)Enum.ToObject(typeof(T), 0);
+
+    /// <inheritdoc />
+    public override bool HandleNull => true;
+
+    public DefaultUnknownEnumConverter()
+    {
+        _enumMemberNamesLazy = new Lazy<Dictionary<T, string>>(
+            () => EnumMemberValues.ToDictionary(x => x.Value, x => x.Key)
+        );
+    }
+
+    /// <inheritdoc />
     public override T Read(
         ref Utf8JsonReader reader,
         Type typeToConvert,
         JsonSerializerOptions options
     )
     {
-        if (reader.TokenType != JsonTokenType.String)
-        {
-            throw new JsonException();
-        }
-
-        var enumText = reader.GetString()?.Replace(" ", "_");
-        if (Enum.TryParse(typeof(T), enumText, true, out var result))
+        if (reader.TokenType is not (JsonTokenType.String or JsonTokenType.PropertyName))
         {
-            return (T)result!;
+            throw new JsonException("Expected String or PropertyName token");
         }
 
-        // Unknown value handling
-        if (Enum.TryParse(typeof(T), "Unknown", true, out var unknownResult))
+        if (reader.GetString() is { } readerString)
         {
-            return (T)unknownResult!;
+            if (EnumMemberValues.TryGetValue(readerString, out var enumMemberValue))
+            {
+                return enumMemberValue;
+            }
         }
 
-        throw new JsonException($"Unable to parse '{enumText}' to enum '{typeof(T)}'.");
+        return UnknownValue;
     }
 
+    /// <inheritdoc />
     public override void Write(Utf8JsonWriter writer, T? value, JsonSerializerOptions options)
     {
         if (value == null)
@@ -41,7 +98,7 @@ public class DefaultUnknownEnumConverter<T> : JsonConverter<T>
             return;
         }
 
-        writer.WriteStringValue(value.GetStringValue().Replace("_", " "));
+        writer.WriteStringValue(EnumMemberNames[value]);
     }
 
     /// <inheritdoc />
@@ -49,41 +106,12 @@ public class DefaultUnknownEnumConverter<T> : JsonConverter<T>
         ref Utf8JsonReader reader,
         Type typeToConvert,
         JsonSerializerOptions options
-    )
-    {
-        if (reader.TokenType != JsonTokenType.PropertyName)
-        {
-            throw new JsonException();
-        }
-
-        var enumText = reader.GetString()?.Replace(" ", "_");
-        if (Enum.TryParse(typeof(T), enumText, true, out var result))
-        {
-            return (T)result!;
-        }
-
-        // Unknown value handling
-        if (Enum.TryParse(typeof(T), "Unknown", true, out var unknownResult))
-        {
-            return (T)unknownResult!;
-        }
-
-        throw new JsonException($"Unable to parse '{enumText}' to enum '{typeof(T)}'.");
-    }
+    ) => Read(ref reader, typeToConvert, options);
 
     /// <inheritdoc />
     public override void WriteAsPropertyName(
         Utf8JsonWriter writer,
         T? value,
         JsonSerializerOptions options
-    )
-    {
-        if (value == null)
-        {
-            writer.WriteNullValue();
-            return;
-        }
-
-        writer.WritePropertyName(value.GetStringValue().Replace("_", " "));
-    }
+    ) => Write(writer, value, options);
 }

StabilityMatrix.Core/Models/Api/CivitFileType.cs

@@ -1,4 +1,5 @@
-using System.Text.Json.Serialization;
+using System.Runtime.Serialization;
+using System.Text.Json.Serialization;
 using StabilityMatrix.Core.Converters.Json;
 
 namespace StabilityMatrix.Core.Models.Api;
@@ -6,8 +7,10 @@ namespace StabilityMatrix.Core.Models.Api;
 [JsonConverter(typeof(DefaultUnknownEnumConverter<CivitFileType>))]
 public enum CivitFileType
 {
+    Unknown,
     Model,
     VAE,
-    Training_Data,
-    Unknown,
+
+    [EnumMember(Value = "Training Data")]
+    TrainingData
 }

StabilityMatrix.Core/Models/Api/CivitModelType.cs

@@ -9,6 +9,8 @@ namespace StabilityMatrix.Core.Models.Api;
 [SuppressMessage("ReSharper", "InconsistentNaming")]
 public enum CivitModelType
 {
+    Unknown,
+
     [ConvertTo<SharedFolderType>(SharedFolderType.StableDiffusion)]
     Checkpoint,
 
@@ -39,6 +41,5 @@ public enum CivitModelType
     Wildcards,
     Workflows,
     Other,
-    All,
-    Unknown
+    All
 }

StabilityMatrix.Tests/Core/DefaultUnknownEnumConverterTests.cs

@@ -0,0 +1,81 @@
+using System.Text.Json;
+using System.Text.Json.Serialization;
+using StabilityMatrix.Core.Converters.Json;
+
+namespace StabilityMatrix.Tests.Core;
+
+[TestClass]
+public class DefaultUnknownEnumConverterTests
+{
+    [TestMethod]
+    [ExpectedException(typeof(JsonException))]
+    public void TestDeserialize_NormalEnum_ShouldError()
+    {
+        const string json = "\"SomeUnknownValue\"";
+
+        JsonSerializer.Deserialize<NormalEnum>(json);
+    }
+
+    [TestMethod]
+    public void TestDeserialize_UnknownEnum_ShouldConvert()
+    {
+        const string json = "\"SomeUnknownValue\"";
+
+        var result = JsonSerializer.Deserialize<UnknownEnum>(json);
+
+        Assert.AreEqual(UnknownEnum.Unknown, result);
+    }
+
+    [TestMethod]
+    public void TestDeserialize_DefaultEnum_ShouldConvert()
+    {
+        const string json = "\"SomeUnknownValue\"";
+
+        var result = JsonSerializer.Deserialize<DefaultEnum>(json);
+
+        Assert.AreEqual(DefaultEnum.CustomDefault, result);
+    }
+
+    [TestMethod]
+    public void TestSerialize_UnknownEnum_ShouldConvert()
+    {
+        const string expected = "\"Unknown\"";
+
+        var result = JsonSerializer.Serialize(UnknownEnum.Unknown);
+
+        Assert.AreEqual(expected, result);
+    }
+
+    [TestMethod]
+    public void TestSerialize_DefaultEnum_ShouldConvert()
+    {
+        const string expected = "\"CustomDefault\"";
+
+        var result = JsonSerializer.Serialize(DefaultEnum.CustomDefault);
+
+        Assert.AreEqual(expected, result);
+    }
+
+    private enum NormalEnum
+    {
+        Unknown,
+        Value1,
+        Value2
+    }
+
+    [JsonConverter(typeof(DefaultUnknownEnumConverter<UnknownEnum>))]
+    private enum UnknownEnum
+    {
+        Unknown,
+        Value1,
+        Value2
+    }
+
+    [JsonConverter(typeof(DefaultUnknownEnumConverter<DefaultEnum>))]
+    private enum DefaultEnum
+    {
+        CustomDefault,
+        Value1,
+        Value2
+    }
+}