feat: add generic support for YamlSerializable, YamlDerivedType and YamlDerivedTypeMapping (#3)

Author: glucaci

src/Yamlify.SourceGenerator/YamlSourceGenerator.cs

@@ -22,6 +22,8 @@ namespace Yamlify.SourceGenerator;
 public sealed class YamlSourceGenerator : IIncrementalGenerator
 {
     private const string YamlSerializableAttribute = "Yamlify.Serialization.YamlSerializableAttribute";
+    private const string YamlSerializableAttributeGeneric = "Yamlify.Serialization.YamlSerializableAttribute<T>";
+    private const string YamlDerivedTypeMappingAttributeGeneric = "Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>";
     private const string YamlSerializerContextBase = "Yamlify.Serialization.YamlSerializerContext";
 
     public void Initialize(IncrementalGeneratorInitializationContext context)
@@ -86,66 +88,137 @@ private static bool IsCandidateClass(SyntaxNode node)
         var ignoreEmptyObjects = false;
         var discriminatorPosition = DiscriminatorPositionMode.PropertyOrder;
 
+        // First pass: collect YamlDerivedTypeMapping attributes
+        // Key: base type display string, Value: list of (discriminator, derivedType)
+        var derivedTypeMappingsFromAttrs = new Dictionary<string, List<(string Discriminator, INamedTypeSymbol DerivedType)>>();
+        
+        foreach (var attributeData in classSymbol.GetAttributes())
+        {
+            var attrOriginalDef = attributeData.AttributeClass?.OriginalDefinition?.ToDisplayString();
+            
+            if (attrOriginalDef == YamlDerivedTypeMappingAttributeGeneric)
+            {
+                // [YamlDerivedTypeMapping<TBase, TDerived>("discriminator")]
+                if (attributeData.AttributeClass is { IsGenericType: true, TypeArguments.Length: 2 } attrClass &&
+                    attrClass.TypeArguments[0] is INamedTypeSymbol mappingBaseType &&
+                    attrClass.TypeArguments[1] is INamedTypeSymbol mappingDerivedType)
+                {
+                    var baseTypeKey = mappingBaseType.ToDisplayString();
+                    
+                    // Get discriminator from constructor argument (optional)
+                    string? discriminator = null;
+                    if (attributeData.ConstructorArguments.Length > 0 && 
+                        attributeData.ConstructorArguments[0].Value is string discValue)
+                    {
+                        discriminator = discValue;
+                    }
+                    discriminator ??= mappingDerivedType.Name;
+                    
+                    if (!derivedTypeMappingsFromAttrs.TryGetValue(baseTypeKey, out var mappings))
+                    {
+                        mappings = new List<(string, INamedTypeSymbol)>();
+                        derivedTypeMappingsFromAttrs[baseTypeKey] = mappings;
+                    }
+                    mappings.Add((discriminator, mappingDerivedType));
+                }
+            }
+        }
+        
+        // Second pass: process YamlSerializable attributes
         foreach (var attributeData in classSymbol.GetAttributes())
         {
             var attrName = attributeData.AttributeClass?.ToDisplayString();
+            var attrOriginalDef = attributeData.AttributeClass?.OriginalDefinition?.ToDisplayString();
+            
+            // Support both [YamlSerializable(typeof(T))] and [YamlSerializable<T>]
+            INamedTypeSymbol? typeArg = null;
+            var isYamlSerializableAttribute = false;
+            
             if (attrName == YamlSerializableAttribute)
             {
+                // Non-generic: [YamlSerializable(typeof(T))]
                 if (attributeData.ConstructorArguments.Length > 0 &&
-                    attributeData.ConstructorArguments[0].Value is INamedTypeSymbol typeArg)
+                    attributeData.ConstructorArguments[0].Value is INamedTypeSymbol ctorArg)
                 {
-                    // Check for per-type PropertyOrdering override
-                    PropertyOrderingMode? typeOrdering = null;
-                    string? typeDiscriminatorPropertyName = null;
-                    List<INamedTypeSymbol>? derivedTypes = null;
-                    List<string>? derivedTypeDiscriminators = null;
-                    
-                    foreach (var namedArg in attributeData.NamedArguments)
+                    typeArg = ctorArg;
+                    isYamlSerializableAttribute = true;
+                }
+            }
+            else if (attrOriginalDef == YamlSerializableAttributeGeneric)
+            {
+                // Generic: [YamlSerializable<T>]
+                if (attributeData.AttributeClass is { IsGenericType: true, TypeArguments.Length: > 0 } attrClass &&
+                    attrClass.TypeArguments[0] is INamedTypeSymbol genericArg)
+                {
+                    typeArg = genericArg;
+                    isYamlSerializableAttribute = true;
+                }
+            }
+            
+            if (isYamlSerializableAttribute && typeArg is not null)
+            {
+                // Check for per-type PropertyOrdering override
+                PropertyOrderingMode? typeOrdering = null;
+                string? typeDiscriminatorPropertyName = null;
+                List<INamedTypeSymbol>? derivedTypes = null;
+                List<string>? derivedTypeDiscriminators = null;
+                
+                foreach (var namedArg in attributeData.NamedArguments)
+                {
+                    if (namedArg.Key == "PropertyOrdering" && namedArg.Value.Value is int orderingValue && orderingValue >= 0)
                     {
-                        if (namedArg.Key == "PropertyOrdering" && namedArg.Value.Value is int orderingValue && orderingValue >= 0)
-                        {
-                            // Only set if not Inherit (-1)
-                            typeOrdering = (PropertyOrderingMode)orderingValue;
-                        }
-                        else if (namedArg.Key == "TypeDiscriminatorPropertyName" && namedArg.Value.Value is string discPropName)
-                        {
-                            typeDiscriminatorPropertyName = discPropName;
-                        }
-                        else if (namedArg.Key == "DerivedTypes" && !namedArg.Value.IsNull)
-                        {
-                            derivedTypes = namedArg.Value.Values
-                                .Where(v => v.Value is INamedTypeSymbol)
-                                .Select(v => (INamedTypeSymbol)v.Value!)
-                                .ToList();
-                        }
-                        else if (namedArg.Key == "DerivedTypeDiscriminators" && !namedArg.Value.IsNull)
-                        {
-                            derivedTypeDiscriminators = namedArg.Value.Values
-                                .Where(v => v.Value is string)
-                                .Select(v => (string)v.Value!)
-                                .ToList();
-                        }
+                        // Only set if not Inherit (-1)
+                        typeOrdering = (PropertyOrderingMode)orderingValue;
                     }
-                    
-                    // Build PolymorphicInfo if polymorphic configuration is specified
-                    PolymorphicInfo? polymorphicConfig = null;
-                    if (typeDiscriminatorPropertyName is not null && derivedTypes is not null && derivedTypes.Count > 0)
+                    else if (namedArg.Key == "TypeDiscriminatorPropertyName" && namedArg.Value.Value is string discPropName)
                     {
-                        var derivedTypeMappings = new List<(string Discriminator, INamedTypeSymbol DerivedType)>();
-                        for (int i = 0; i < derivedTypes.Count; i++)
-                        {
-                            var derivedType = derivedTypes[i];
-                            // Use explicit discriminator if provided, otherwise use type name
-                            var discriminator = (derivedTypeDiscriminators is not null && i < derivedTypeDiscriminators.Count)
-                                ? derivedTypeDiscriminators[i]
-                                : derivedType.Name;
-                            derivedTypeMappings.Add((discriminator, derivedType));
-                        }
-                        polymorphicConfig = new PolymorphicInfo(typeDiscriminatorPropertyName, derivedTypeMappings);
+                        typeDiscriminatorPropertyName = discPropName;
                     }
-                    
-                    typesToGenerate.Add(new TypeToGenerate(typeArg, typeOrdering, polymorphicConfig));
+                    else if (namedArg.Key == "DerivedTypes" && !namedArg.Value.IsNull)
+                    {
+                        derivedTypes = namedArg.Value.Values
+                            .Where(v => v.Value is INamedTypeSymbol)
+                            .Select(v => (INamedTypeSymbol)v.Value!)
+                            .ToList();
+                    }
+                    else if (namedArg.Key == "DerivedTypeDiscriminators" && !namedArg.Value.IsNull)
+                    {
+                        derivedTypeDiscriminators = namedArg.Value.Values
+                            .Where(v => v.Value is string)
+                            .Select(v => (string)v.Value!)
+                            .ToList();
+                    }
+                }
+                
+                // Build PolymorphicInfo if polymorphic configuration is specified
+                PolymorphicInfo? polymorphicConfig = null;
+                var typeKey = typeArg.ToDisplayString();
+                
+                // Check for derived type mappings from YamlDerivedTypeMappingAttribute
+                if (typeDiscriminatorPropertyName is not null && 
+                    derivedTypeMappingsFromAttrs.TryGetValue(typeKey, out var mappingsFromAttr) && 
+                    mappingsFromAttr.Count > 0)
+                {
+                    // Use mappings from YamlDerivedTypeMappingAttribute
+                    polymorphicConfig = new PolymorphicInfo(typeDiscriminatorPropertyName, mappingsFromAttr);
                 }
+                else if (typeDiscriminatorPropertyName is not null && derivedTypes is not null && derivedTypes.Count > 0)
+                {
+                    // Use inline DerivedTypes/DerivedTypeDiscriminators arrays
+                    var derivedTypeMappings = new List<(string Discriminator, INamedTypeSymbol DerivedType)>();
+                    for (int i = 0; i < derivedTypes.Count; i++)
+                    {
+                        var derivedType = derivedTypes[i];
+                        // Use explicit discriminator if provided, otherwise use type name
+                        var discriminator = (derivedTypeDiscriminators is not null && i < derivedTypeDiscriminators.Count)
+                            ? derivedTypeDiscriminators[i]
+                            : derivedType.Name;
+                        derivedTypeMappings.Add((discriminator, derivedType));
+                    }
+                    polymorphicConfig = new PolymorphicInfo(typeDiscriminatorPropertyName, derivedTypeMappings);
+                }
+                
+                typesToGenerate.Add(new TypeToGenerate(typeArg, typeOrdering, polymorphicConfig));
             }
             else if (attrName == "Yamlify.Serialization.YamlSourceGenerationOptionsAttribute")
             {

src/Yamlify/PublicAPI.Unshipped.txt

@@ -1 +1,10 @@
-#nullable enable
+#nullable enable
+Yamlify.Serialization.YamlDerivedTypeAttribute<T>
+Yamlify.Serialization.YamlDerivedTypeAttribute<T>.YamlDerivedTypeAttribute(string? typeDiscriminator = null) -> void
+Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>
+Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>.BaseType.get -> System.Type!
+Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>.DerivedType.get -> System.Type!
+Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>.TypeDiscriminator.get -> string?
+Yamlify.Serialization.YamlDerivedTypeMappingAttribute<TBase, TDerived>.YamlDerivedTypeMappingAttribute(string? typeDiscriminator = null) -> void
+Yamlify.Serialization.YamlSerializableAttribute<T>
+Yamlify.Serialization.YamlSerializableAttribute<T>.YamlSerializableAttribute() -> void

src/Yamlify/Serialization/YamlDerivedTypeAttribute.cs

@@ -4,7 +4,7 @@ namespace Yamlify.Serialization;
 /// Specifies polymorphic type information.
 /// </summary>
 [AttributeUsage(AttributeTargets.Class | AttributeTargets.Interface, AllowMultiple = true)]
-public sealed class YamlDerivedTypeAttribute : Attribute
+public class YamlDerivedTypeAttribute : Attribute
 {
     /// <summary>
     /// Gets the derived type.

src/Yamlify/Serialization/YamlDerivedTypeAttributeOfT.cs

@@ -0,0 +1,25 @@
+namespace Yamlify.Serialization;
+
+/// <summary>
+/// Generic version of <see cref="YamlDerivedTypeAttribute"/> for a more type-safe API.
+/// </summary>
+/// <typeparam name="T">The derived type.</typeparam>
+/// <example>
+/// <code>
+/// [YamlPolymorphic(TypeDiscriminatorPropertyName = "type")]
+/// [YamlDerivedType&lt;Dog&gt;("dog")]
+/// [YamlDerivedType&lt;Cat&gt;("cat")]
+/// public abstract class Animal { }
+/// </code>
+/// </example>
+[AttributeUsage(AttributeTargets.Class | AttributeTargets.Interface, AllowMultiple = true)]
+public sealed class YamlDerivedTypeAttribute<T> : YamlDerivedTypeAttribute
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="YamlDerivedTypeAttribute{T}"/> class.
+    /// </summary>
+    /// <param name="typeDiscriminator">The type discriminator value. If null, the type name is used.</param>
+    public YamlDerivedTypeAttribute(string? typeDiscriminator = null) : base(typeof(T), typeDiscriminator)
+    {
+    }
+}

src/Yamlify/Serialization/YamlDerivedTypeMappingAttribute.cs

@@ -0,0 +1,58 @@
+namespace Yamlify.Serialization;
+
+/// <summary>
+/// Specifies a polymorphic type mapping between a base type and a derived type for the serializer context.
+/// This attribute provides a type-safe way to declare derived type mappings on the context class.
+/// </summary>
+/// <typeparam name="TBase">The base type or interface for the polymorphic hierarchy.</typeparam>
+/// <typeparam name="TDerived">The derived type that implements or extends the base type.</typeparam>
+/// <remarks>
+/// <para>
+/// Use this attribute together with <see cref="YamlSerializableAttribute{T}"/> to configure
+/// polymorphic serialization directly on the context class.
+/// </para>
+/// <para>
+/// When using this attribute, you must also register the base type with 
+/// <see cref="YamlSerializableAttribute{T}"/> and set its <see cref="YamlSerializableAttribute.TypeDiscriminatorPropertyName"/>.
+/// </para>
+/// </remarks>
+/// <example>
+/// <code>
+/// [YamlSerializable&lt;IAnimal&gt;(TypeDiscriminatorPropertyName = "type")]
+/// [YamlDerivedTypeMapping&lt;IAnimal, Dog&gt;("dog")]
+/// [YamlDerivedTypeMapping&lt;IAnimal, Cat&gt;("cat")]
+/// [YamlSerializable&lt;Dog&gt;]
+/// [YamlSerializable&lt;Cat&gt;]
+/// public partial class MyContext : YamlSerializerContext { }
+/// </code>
+/// </example>
+[AttributeUsage(AttributeTargets.Class, AllowMultiple = true)]
+public sealed class YamlDerivedTypeMappingAttribute<TBase, TDerived> : Attribute
+    where TDerived : TBase
+{
+    /// <summary>
+    /// Gets the base type for the polymorphic hierarchy.
+    /// </summary>
+    public Type BaseType => typeof(TBase);
+
+    /// <summary>
+    /// Gets the derived type.
+    /// </summary>
+    public Type DerivedType => typeof(TDerived);
+
+    /// <summary>
+    /// Gets the type discriminator value used in YAML to identify this derived type.
+    /// </summary>
+    public string? TypeDiscriminator { get; }
+
+    /// <summary>
+    /// Initializes a new instance of the <see cref="YamlDerivedTypeMappingAttribute{TBase, TDerived}"/> class.
+    /// </summary>
+    /// <param name="typeDiscriminator">
+    /// The type discriminator value. If null, the type name of <typeparamref name="TDerived"/> is used.
+    /// </param>
+    public YamlDerivedTypeMappingAttribute(string? typeDiscriminator = null)
+    {
+        TypeDiscriminator = typeDiscriminator;
+    }
+}

src/Yamlify/Serialization/YamlSerializableAttribute.cs

@@ -31,7 +31,7 @@ namespace Yamlify.Serialization;
 /// </code>
 /// </example>
 [AttributeUsage(AttributeTargets.Class, AllowMultiple = true)]
-public sealed class YamlSerializableAttribute : Attribute
+public class YamlSerializableAttribute : Attribute
 {
     /// <summary>
     /// Gets the type for which to generate serialization metadata.

src/Yamlify/Serialization/YamlSerializableAttributeOfT.cs

@@ -0,0 +1,27 @@
+namespace Yamlify.Serialization;
+
+/// <summary>
+/// Generic version of <see cref="YamlSerializableAttribute"/> for a more type-safe API.
+/// </summary>
+/// <typeparam name="T">The type for which to generate serialization metadata.</typeparam>
+/// <example>
+/// <code>
+/// [YamlSerializable&lt;Person&gt;]
+/// [YamlSerializable&lt;Address&gt;]
+/// [YamlSerializable&lt;IAnimal&gt;(
+///     TypeDiscriminatorPropertyName = "type",
+///     DerivedTypes = new[] { typeof(Dog), typeof(Cat) },
+///     DerivedTypeDiscriminators = new[] { "dog", "cat" })]
+/// public partial class MyContext : YamlSerializerContext { }
+/// </code>
+/// </example>
+[AttributeUsage(AttributeTargets.Class, AllowMultiple = true)]
+public sealed class YamlSerializableAttribute<T> : YamlSerializableAttribute
+{
+    /// <summary>
+    /// Initializes a new instance of the <see cref="YamlSerializableAttribute{T}"/> class.
+    /// </summary>
+    public YamlSerializableAttribute() : base(typeof(T))
+    {
+    }
+}