feat: Resolve additional CustomHandler type arguments from matching type (#39)

Author: Dreamescaper

README.md

@@ -133,6 +133,24 @@ public static partial class ServiceCollectionExtensions
 }
 ```
 
+### Apply EF Core IEntityTypeConfiguration types
+
+```csharp
+public static partial class ModelBuilderExtensions
+{
+    [GenerateServiceRegistrations(AssignableTo = typeof(IEntityTypeConfiguration<>), CustomHandler = nameof(ApplyConfiguration))]
+    public static partial ModelBuilder ApplyEntityConfigurations(this ModelBuilder modelBuilder);
+
+    private static void ApplyConfiguration<T, TEntity>(ModelBuilder modelBuilder)
+        where T : IEntityTypeConfiguration<TEntity>, new()
+        where TEntity : class
+    {
+        modelBuilder.ApplyConfiguration(new T());
+    }
+}
+```
+
+
 
 ## Parameters
 

ServiceScan.SourceGenerator.Tests/CustomHandlerTests.cs

@@ -268,7 +268,6 @@ public static partial void ProcessServices()
         Assert.Equal(expected, results.GeneratedTrees[1].ToString());
     }
 
-
     [Fact]
     public void AddMultipleCustomHandlerAttributesWithSameCustomHandler()
     {
@@ -319,6 +318,85 @@ public static partial void ProcessServices()
         Assert.Equal(expected, results.GeneratedTrees[1].ToString());
     }
 
+    [Fact]
+    public void ResolveCustomHandlerGenericArguments()
+    {
+        var source = $$"""
+            using ServiceScan.SourceGenerator;
+            
+            namespace GeneratorTests;
+                    
+            public static partial class ModelBuilderExtensions
+            {
+                [GenerateServiceRegistrations(AssignableTo = typeof(IEntityTypeConfiguration<>), CustomHandler = nameof(ApplyConfiguration))]
+                public static partial ModelBuilder ApplyEntityConfigurations(this ModelBuilder modelBuilder);
+
+                private static void ApplyConfiguration<T, TEntity>(ModelBuilder modelBuilder)
+                    where T : IEntityTypeConfiguration<TEntity>, new()
+                    where TEntity : class
+                {
+                    modelBuilder.ApplyConfiguration(new T());
+                }
+            }
+            """;
+
+        var infra = """
+            public interface IEntityTypeConfiguration<TEntity> where TEntity : class
+            {
+                void Configure(EntityTypeBuilder<TEntity> builder);
+            }
+
+            public class EntityTypeBuilder<TEntity> where TEntity : class;
+
+            public class ModelBuilder
+            {
+                public ModelBuilder ApplyConfiguration<TEntity>(IEntityTypeConfiguration<TEntity> configuration) where TEntity : class
+                {
+                    return this;
+                }
+            }
+            """;
+
+        var configurations = """
+            namespace GeneratorTests;
+            
+            public class EntityA;
+            public class EntityB;
+
+            public class EntityAConfiguration : IEntityTypeConfiguration<EntityA>
+            {
+                public void Configure(EntityTypeBuilder<EntityA> builder) { }
+            }
+
+            public class EntityBConfiguration : IEntityTypeConfiguration<EntityB>
+            {
+                public void Configure(EntityTypeBuilder<EntityB> builder) { }
+            }
+            """;
+
+        var compilation = CreateCompilation(source, infra, configurations);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var expected = $$"""
+            namespace GeneratorTests;
+
+            public static partial class ModelBuilderExtensions
+            {
+                public static partial global::ModelBuilder ApplyEntityConfigurations(this global::ModelBuilder modelBuilder)
+                {
+                    ApplyConfiguration<global::GeneratorTests.EntityAConfiguration, global::GeneratorTests.EntityA>(modelBuilder);
+                    ApplyConfiguration<global::GeneratorTests.EntityBConfiguration, global::GeneratorTests.EntityB>(modelBuilder);
+                    return modelBuilder;
+                }
+            }
+            """;
+        Assert.Equal(expected, results.GeneratedTrees[1].ToString());
+    }
+
 
     private static Compilation CreateCompilation(params string[] source)
     {

ServiceScan.SourceGenerator/DependencyInjectionGenerator.FindServicesToRegister.cs

@@ -37,13 +37,31 @@ private static DiagnosticModel<MethodImplementationModel> FindServicesToRegister
 
                 if (attribute.CustomHandler != null)
                 {
-                    customHandlers.Add(new CustomHandlerModel(attribute.CustomHandler, implementationType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)));
+                    // If CustomHandler method has multiple type parameters, which are resolvable from the first one - we try to provide them.
+                    // e.g. ApplyConfiguration<T, TEntity>(ModelBuilder modelBuilder) where T : IEntityTypeConfiguration<TEntity>
+                    if (attribute.CustomHandlerTypeParametersCount > 1 && matchedTypes != null)
+                    {
+                        foreach (var matchedType in matchedTypes)
+                        {
+                            EquatableArray<string> typeArguments =
+                                [
+                                    implementationType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat),
+                                    .. matchedType.TypeArguments.Select(a => a.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat))
+                                ];
+
+                            customHandlers.Add(new CustomHandlerModel(attribute.CustomHandler, typeArguments));
+                        }
+                    }
+                    else
+                    {
+                        customHandlers.Add(new CustomHandlerModel(attribute.CustomHandler, [implementationType.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat)]));
+                    }
                 }
                 else
                 {
                     var serviceTypes = (attribute.AsSelf, attribute.AsImplementedInterfaces) switch
                     {
-                        (true, true) => new[] { implementationType }.Concat(GetSuitableInterfaces(implementationType)),
+                        (true, true) => [implementationType, .. GetSuitableInterfaces(implementationType)],
                         (false, true) => GetSuitableInterfaces(implementationType),
                         (true, false) => [implementationType],
                         _ => matchedTypes ?? [implementationType]

ServiceScan.SourceGenerator/DependencyInjectionGenerator.ParseMethodModel.cs

@@ -69,6 +69,15 @@ public partial class DependencyInjectionGenerator
 
                 if (!typesMatch)
                     return Diagnostic.Create(CustomHandlerMethodHasIncorrectSignature, attribute.Location);
+
+                // If CustomHandler has more than 1 type parameters, we try to resolve them from 
+                // matched assignableTo type arguments.
+                // e.g. ApplyConfiguration<T, TEntity>(ModelBuilder modelBuilder) where T : IEntityTypeConfiguration<TEntity>
+                if (customHandlerMethod.TypeParameters.Length > 1
+                    && customHandlerMethod.TypeParameters.Length != attribute.AssignableToTypeParametersCount + 1)
+                {
+                    return Diagnostic.Create(CustomHandlerMethodHasIncorrectSignature, attribute.Location);
+                }
             }
 
             if (attributeData[i].HasErrors)

ServiceScan.SourceGenerator/DependencyInjectionGenerator.cs

@@ -108,7 +108,11 @@ private static string GenerateRegistrationsSource(MethodModel method, EquatableA
     private static string GenerateCustomHandlingSource(MethodModel method, EquatableArray<CustomHandlerModel> customHandlers)
     {
         var invocations = string.Join("\n", customHandlers.Select(h =>
-            $"        {h.HandlerMethodName}<{h.TypeName}>({string.Join(", ", method.Parameters.Select(p => p.Name))});"));
+        {
+            var genericArguments = string.Join(", ", h.TypeArguments);
+            var arguments = string.Join(", ", method.Parameters.Select(p => p.Name));
+            return $"        {h.HandlerMethodName}<{genericArguments}>({arguments});";
+        }));
 
         var namespaceDeclaration = method.Namespace is null ? "" : $"namespace {method.Namespace};";
         var parameters = string.Join(",", method.Parameters.Select((p, i) =>

ServiceScan.SourceGenerator/Model/AttributeModel.cs

@@ -7,6 +7,7 @@ enum KeySelectorType { Method, GenericMethod, TypeMember };
 
 record AttributeModel(
     string? AssignableToTypeName,
+    int AssignableToTypeParametersCount,
     string? AssemblyNameFilter,
     EquatableArray<string>? AssignableToGenericArguments,
     string? AssemblyOfTypeName,
@@ -20,6 +21,7 @@ record AttributeModel(
     string? KeySelector,
     KeySelectorType? KeySelectorType,
     string? CustomHandler,
+    int CustomHandlerTypeParametersCount,
     bool AsImplementedInterfaces,
     bool AsSelf,
     Location Location,
@@ -42,6 +44,8 @@ public static AttributeModel Create(AttributeData attribute, IMethodSymbol metho
         var keySelector = attribute.NamedArguments.FirstOrDefault(a => a.Key == "KeySelector").Value.Value as string;
         var customHandler = attribute.NamedArguments.FirstOrDefault(a => a.Key == "CustomHandler").Value.Value as string;
 
+        var assignableToTypeParametersCount = assignableTo?.TypeParameters.Length ?? 0;
+
         KeySelectorType? keySelectorType = null;
         if (keySelector != null)
         {
@@ -59,6 +63,16 @@ public static AttributeModel Create(AttributeData attribute, IMethodSymbol metho
             }
         }
 
+        var customHandlerGenericParameters = 0;
+        if (customHandler != null)
+        {
+            var customHandlerMethod = method.ContainingType.GetMembers()
+                .OfType<IMethodSymbol>()
+                .FirstOrDefault(m => m.IsStatic && m.Name == customHandler);
+
+            customHandlerGenericParameters = customHandlerMethod?.TypeParameters.Length ?? 0;
+        }
+
         if (string.IsNullOrWhiteSpace(typeNameFilter))
             typeNameFilter = null;
 
@@ -97,6 +111,7 @@ public static AttributeModel Create(AttributeData attribute, IMethodSymbol metho
 
         return new(
             assignableToTypeName,
+            assignableToTypeParametersCount,
             assemblyNameFilter,
             assignableToGenericArguments,
             assemblyOfTypeName,
@@ -110,6 +125,7 @@ public static AttributeModel Create(AttributeData attribute, IMethodSymbol metho
             keySelector,
             keySelectorType,
             customHandler,
+            customHandlerGenericParameters,
             asImplementedInterfaces,
             asSelf,
             location,

ServiceScan.SourceGenerator/Model/ServiceRegistrationModel.cs

@@ -11,4 +11,4 @@ record ServiceRegistrationModel(
 
 record CustomHandlerModel(
     string HandlerMethodName,
-    string TypeName);
+    EquatableArray<string> TypeArguments);

version.json

@@ -1,6 +1,6 @@
 {
   "$schema": "https://raw.githubusercontent.com/dotnet/Nerdbank.GitVersioning/main/src/NerdBank.GitVersioning/version.schema.json",
-  "version": "2.1",
+  "version": "2.2",
   "publicReleaseRefSpec": [
     "^refs/heads/main",
     "^refs/heads/v\\d+(?:\\.\\d+)?$"