Allow multiple CustomHandler attributes on the same method (#29)

Author: Dreamescaper

ServiceScan.SourceGenerator.Tests/CustomHandlerTests.cs

@@ -217,6 +217,109 @@ public static partial IServiceCollection ProcessServices(this IServiceCollection
         Assert.Equal(expected, results.GeneratedTrees[1].ToString());
     }
 
+    [Fact]
+    public void AddMultipleCustomHandlerAttributesWithDifferentCustomHandler()
+    {
+        var source = $$"""
+            using ServiceScan.SourceGenerator;
+            
+            namespace GeneratorTests;
+                    
+            public static partial class ServicesExtensions
+            {
+                [GenerateServiceRegistrations(AssignableTo = typeof(IFirstService), CustomHandler = nameof(HandleFirstType))]
+                [GenerateServiceRegistrations(AssignableTo = typeof(ISecondService), CustomHandler = nameof(HandleSecondType))]
+                public static partial void ProcessServices();
+
+                private static void HandleFirstType<T>() => System.Console.WriteLine("First:" + typeof(T).Name);
+                private static void HandleSecondType<T>() => System.Console.WriteLine("Second:" + typeof(T).Name);
+            }
+            """;
+
+        var services =
+            """
+            namespace GeneratorTests;
+
+            public interface IFirstService { }
+            public interface ISecondService { }
+            public class MyService1 : IFirstService { }
+            public class MyService2 : ISecondService { }
+            """;
+
+        var compilation = CreateCompilation(source, services);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var expected = $$"""
+            namespace GeneratorTests;
+
+            public static partial class ServicesExtensions
+            {
+                public static partial void ProcessServices()
+                {
+                    HandleFirstType<global::GeneratorTests.MyService1>();
+                    HandleSecondType<global::GeneratorTests.MyService2>();
+                }
+            }
+            """;
+        Assert.Equal(expected, results.GeneratedTrees[1].ToString());
+    }
+
+
+    [Fact]
+    public void AddMultipleCustomHandlerAttributesWithSameCustomHandler()
+    {
+        var source = $$"""
+            using ServiceScan.SourceGenerator;
+            
+            namespace GeneratorTests;
+                    
+            public static partial class ServicesExtensions
+            {
+                [GenerateServiceRegistrations(AssignableTo = typeof(IFirstService), CustomHandler = nameof(HandleType))]
+                [GenerateServiceRegistrations(AssignableTo = typeof(ISecondService), CustomHandler = nameof(HandleType))]
+                public static partial void ProcessServices();
+
+                private static void HandleType<T>() => System.Console.WriteLine(typeof(T).Name);
+            }
+            """;
+
+        var services =
+            """
+            namespace GeneratorTests;
+
+            public interface IFirstService { }
+            public interface ISecondService { }
+            public class MyService1 : IFirstService { }
+            public class MyService2 : ISecondService { }
+            """;
+
+        var compilation = CreateCompilation(source, services);
+
+        var results = CSharpGeneratorDriver
+            .Create(_generator)
+            .RunGenerators(compilation)
+            .GetRunResult();
+
+        var expected = $$"""
+            namespace GeneratorTests;
+
+            public static partial class ServicesExtensions
+            {
+                public static partial void ProcessServices()
+                {
+                    HandleType<global::GeneratorTests.MyService1>();
+                    HandleType<global::GeneratorTests.MyService2>();
+                }
+            }
+            """;
+        Assert.Equal(expected, results.GeneratedTrees[1].ToString());
+    }
+
+
     private static Compilation CreateCompilation(params string[] source)
     {
         var path = Path.GetDirectoryName(typeof(object).Assembly.Location)!;

ServiceScan.SourceGenerator/DependencyInjectionGenerator.ParseMethodModel.cs

@@ -16,8 +16,9 @@ public partial class DependencyInjectionGenerator
         if (!method.IsPartialDefinition)
             return Diagnostic.Create(NotPartialDefinition, method.Locations[0]);
 
-        var hasCustomHandler = false;
-        var attributeData = new AttributeModel[context.Attributes.Length];
+        var attributeData = context.Attributes.Select(a => AttributeModel.Create(a, method)).ToArray();
+        var hasCustomHandlers = attributeData.Any(a => a.CustomHandler != null);
+
         for (var i = 0; i < context.Attributes.Length; i++)
         {
             var attribute = AttributeModel.Create(context.Attributes[i], method);
@@ -26,9 +27,8 @@ public partial class DependencyInjectionGenerator
             if (!attribute.HasSearchCriteria)
                 return Diagnostic.Create(MissingSearchCriteria, attribute.Location);
 
-            hasCustomHandler |= attribute.CustomHandler != null;
-            if (hasCustomHandler && context.Attributes.Length != 1)
-                return Diagnostic.Create(OnlyOneCustomHandlerAllowed, attribute.Location);
+            if (hasCustomHandlers && attribute.CustomHandler == null)
+                return Diagnostic.Create(CantMixRegularAndCustomHandlerRegistrations, attribute.Location);
 
             if (attribute.KeySelector != null)
             {
@@ -72,7 +72,7 @@ public partial class DependencyInjectionGenerator
                 return null;
         }
 
-        if (!hasCustomHandler)
+        if (!hasCustomHandlers)
         {
             var serviceCollectionType = context.SemanticModel.Compilation.GetTypeByMetadataName("Microsoft.Extensions.DependencyInjection.IServiceCollection");
 

ServiceScan.SourceGenerator/DiagnosticDescriptors.cs

@@ -46,9 +46,9 @@ public static class DiagnosticDescriptors
         DiagnosticSeverity.Error,
         true);
 
-    public static readonly DiagnosticDescriptor OnlyOneCustomHandlerAllowed = new("DI0008",
-        "Only one GenerateServiceRegistrations attribute is allowed when CustomHandler used",
-        "Only one GenerateServiceRegistrations attribute is allowed when CustomHandler used",
+    public static readonly DiagnosticDescriptor CantMixRegularAndCustomHandlerRegistrations = new("DI0008",
+        "It's not allowed to mix GenerateServiceRegistrations attributes with and without CustomHandler on the same method",
+        "It's not allowed to mix GenerateServiceRegistrations attributes with and without CustomHandler on the same method",
         "Usage",
         DiagnosticSeverity.Error,
         true);
@@ -69,7 +69,7 @@ public static class DiagnosticDescriptors
 
     public static readonly DiagnosticDescriptor CustomHandlerMethodHasIncorrectSignature = new("DI0011",
         "Provided CustomHandler method has incorrect signature",
-        "CustomHandler method must be generic, and must have the same parameters as the method with an attribute",
+        "CustomHandler method must be generic, and must have the same parameters as the method with the attribute",
         "Usage",
         DiagnosticSeverity.Error,
         true);