Add support for multiple keyed service registrations

Fixes #108
devlooped · Nov 11, 2024 · e1adab8 · e1adab8
1 parent 8ac296b
commit e1adab8
Show file tree

Hide file tree

Showing 5 changed files with 94 additions and 48 deletions.
diff --git a/readme.md b/readme.md
@@ -81,6 +81,7 @@ public class SmsNotificationService : INotificationService
 }
 
 [Service<string>("email")]
+[Service<string>("default")]
 public class EmailNotificationService : INotificationService
 {
     public string Notify(string message) => $"[Email] {message}";
@@ -101,6 +102,8 @@ public class SmsService([FromKeyedServices("sms")] INotificationService sms)
 In this case, when resolving the `SmsService` from the service provider, the 
 right `INotificationService` will be injected, based on the key provided.
 
+Note you can also register the same service using multiple keys, as shown in the 
+`EmailNotificationService` above.
 
 ## How It Works
 

diff --git a/src/DependencyInjection.Attributed.Tests/GenerationTests.cs b/src/DependencyInjection.Attributed.Tests/GenerationTests.cs
@@ -200,6 +200,23 @@ public void ResolvesKeyedFromContracts()
         Assert.NotNull(singleton.Dependency);
     }
 
+    [Fact]
+    public void ResolveMultipleKeys()
+    {
+        var collection = new ServiceCollection();
+        collection.AddServices();
+        var services = collection.BuildServiceProvider();
+
+        var sms = services.GetRequiredKeyedService<INotificationService>("sms");
+        var email = services.GetRequiredKeyedService<INotificationService>("email");
+        var def = services.GetRequiredKeyedService<INotificationService>("default");
+
+        // Each gets its own instance, since we can't tell apart. Lifetimes can also be disparate.
+        Assert.NotSame(sms, email);
+        Assert.NotSame(sms, def);
+        Assert.NotSame(email, def);
+    }
+
     [Fact]
     public void ResolvesDependency()
     {
@@ -349,4 +366,22 @@ public class DependencyFromKeyedContract([Import("contract")] KeyedByContractNam
 
 public interface IService { }
 [Service]
-class InternalService : IService { }
+class InternalService : IService { }
+
+public interface INotificationService
+{
+    string Notify(string message);
+}
+
+[Service<string>("sms")]
+public class SmsNotificationService : INotificationService
+{
+    public string Notify(string message) => $"[SMS] {message}";
+}
+
+[Service<string>("email")]
+[Service<string>("default")]
+public class EmailNotificationService : INotificationService
+{
+    public string Notify(string message) => $"[Email] {message}";
+}
diff --git a/src/DependencyInjection.Attributed/Attributed.csproj b/src/DependencyInjection.Attributed/Attributed.csproj
@@ -16,16 +16,22 @@
   <ItemGroup>
     <None Update="Devlooped.Extensions.DependencyInjection.Attributed.props" CopyToOutputDirectory="PreserveNewest" PackFolder="buildTransitive" />
     <None Update="Devlooped.Extensions.DependencyInjection.Attributed.targets" CopyToOutputDirectory="PreserveNewest" PackFolder="buildTransitive" />
-    <EmbeddedResource Include="ServiceAttribute*.cs;AddServicesExtension.cs" />
   </ItemGroup>
 
   <ItemGroup>
     <PackageReference Include="NuGetizer" Version="1.2.1" />
     <PackageReference Include="Microsoft.CodeAnalysis.CSharp" Version="4.2.0" Pack="false" />
     <PackageReference Include="Microsoft.Extensions.DependencyInjection.Abstractions" Version="6.0.0" />
+    <PackageReference Include="PolySharp" Version="1.14.1" PrivateAssets="all" />
     <PackageReference Include="ThisAssembly.Resources" Version="2.0.8" PrivateAssets="all" />
   </ItemGroup>
 
+  <Target Name="AddEmbeddedResources" BeforeTargets="SplitResourcesByCulture" Condition="'$(DesignTimeBuild)' != 'true'">
+    <ItemGroup>
+      <EmbeddedResource Include="ServiceAttribute*.cs;AddServicesExtension.cs" Type="Non-Resx" />
+    </ItemGroup>
+  </Target>
+
   <Target Name="PokePackageVersion" BeforeTargets="GetPackageContents" DependsOnTargets="CopyFilesToOutputDirectory" Condition="'$(dotnet-nugetize)' == '' and Exists('$(OutputPath)\Devlooped.Extensions.DependencyInjection.Attributed.props')">
     <XmlPoke XmlInputPath="$(OutputPath)\Devlooped.Extensions.DependencyInjection.Attributed.props" Query="/Project/PropertyGroup/DevloopedExtensionsDependencyInjectionVersion" Value="$(PackageVersion)" />
   </Target>

diff --git a/src/DependencyInjection.Attributed/IncrementalGenerator.cs b/src/DependencyInjection.Attributed/IncrementalGenerator.cs
@@ -18,6 +18,8 @@ namespace Devlooped.Extensions.DependencyInjection.Attributed;
 [Generator(LanguageNames.CSharp)]
 public class IncrementalGenerator : IIncrementalGenerator
 {
+    record ServiceSymbol(INamedTypeSymbol Type, TypedConstant? Key, int Lifetime);
+
     public void Initialize(IncrementalGeneratorInitializationContext context)
     {
         var types = context.CompilationProvider.SelectMany((x, c) =>
@@ -58,65 +60,65 @@ bool IsExport(AttributeData attr)
         // NOTE: we recognize the attribute by name, not precise type. This makes the generator 
         // more flexible and avoids requiring any sort of run-time dependency.
         var services = types
-            .Select((x, _) =>
+            .SelectMany((x, _) =>
             {
                 var name = x.Name;
                 var attrs = x.GetAttributes();
-                var serviceAttr = attrs.FirstOrDefault(IsService) ?? attrs.FirstOrDefault(IsKeyedService);
-                var service = serviceAttr != null || attrs.Any(IsExport);
+                var services = new List<ServiceSymbol>();
 
-                if (!service)
-                    return null;
+                foreach (var attr in attrs)
+                {
+                    var serviceAttr = IsService(attr) || IsKeyedService(attr) ? attr : null;
+                    if (serviceAttr == null && !IsExport(attr))
+                        continue;
 
-                TypedConstant? key = default;
+                    TypedConstant? key = default;
 
-                // Default lifetime is singleton for [Service], Transient for MEF
-                var lifetime = serviceAttr != null ? 0 : 2;
-                if (serviceAttr != null)
-                {
-                    if (IsKeyedService(serviceAttr))
+                    // Default lifetime is singleton for [Service], Transient for MEF
+                    var lifetime = serviceAttr != null ? 0 : 2;
+                    if (serviceAttr != null)
                     {
-                        key = serviceAttr.ConstructorArguments[0];
-                        lifetime = (int)serviceAttr.ConstructorArguments[1].Value!;
+                        if (IsKeyedService(serviceAttr))
+                        {
+                            key = serviceAttr.ConstructorArguments[0];
+                            lifetime = (int)serviceAttr.ConstructorArguments[1].Value!;
+                        }
+                        else
+                        {
+                            lifetime = (int)serviceAttr.ConstructorArguments[0].Value!;
+                        }
                     }
                     else
                     {
-                        lifetime = (int)serviceAttr.ConstructorArguments[0].Value!;
-                    }
-                }
-                else
-                {
-                    // In NuGet MEF, [Shared] makes exports singleton
-                    if (attrs.Any(a => a.AttributeClass?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.Composition.SharedAttribute"))
-                    {
-                        lifetime = 0;
-                    }
-                    // In .NET MEF, [PartCreationPolicy(CreationPolicy.Shared)] does it.
-                    else if (attrs.Any(a =>
-                        a.AttributeClass?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.ComponentModel.Composition.PartCreationPolicyAttribute" &&
-                        a.ConstructorArguments.Length == 1 &&
-                        a.ConstructorArguments[0].Kind == TypedConstantKind.Enum &&
-                        a.ConstructorArguments[0].Type?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.ComponentModel.Composition.CreationPolicy" &&
-                        (int)a.ConstructorArguments[0].Value! == 1))
-                    {
-                        lifetime = 0;
-                    }
+                        // In NuGet MEF, [Shared] makes exports singleton
+                        if (attrs.Any(a => a.AttributeClass?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.Composition.SharedAttribute"))
+                        {
+                            lifetime = 0;
+                        }
+                        // In .NET MEF, [PartCreationPolicy(CreationPolicy.Shared)] does it.
+                        else if (attrs.Any(a =>
+                            a.AttributeClass?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.ComponentModel.Composition.PartCreationPolicyAttribute" &&
+                            a.ConstructorArguments.Length == 1 &&
+                            a.ConstructorArguments[0].Kind == TypedConstantKind.Enum &&
+                            a.ConstructorArguments[0].Type?.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat) == "global::System.ComponentModel.Composition.CreationPolicy" &&
+                            (int)a.ConstructorArguments[0].Value! == 1))
+                        {
+                            lifetime = 0;
+                        }
 
-                    // Consider the [Export(contractName)] as a keyed service with the contract name as the key.
-                    if (attrs.FirstOrDefault(IsExport) is { } export &&
-                        export.ConstructorArguments.Length > 0 &&
-                        export.ConstructorArguments[0].Kind == TypedConstantKind.Primitive)
-                    {
-                        key = export.ConstructorArguments[0];
+                        // Consider the [Export(contractName)] as a keyed service with the contract name as the key.
+                        if (attrs.FirstOrDefault(IsExport) is { } export &&
+                            export.ConstructorArguments.Length > 0 &&
+                            export.ConstructorArguments[0].Kind == TypedConstantKind.Primitive)
+                        {
+                            key = export.ConstructorArguments[0];
+                        }
                     }
+
+                    services.Add(new(x, key, lifetime));
                 }
 
-                return new
-                {
-                    Type = x,
-                    Key = key,
-                    Lifetime = lifetime
-                };
+                return services.ToImmutableArray();
             })
             .Where(x => x != null);
 

diff --git a/src/DependencyInjection.Attributed/ServiceAttribute`1.cs b/src/DependencyInjection.Attributed/ServiceAttribute`1.cs
@@ -9,7 +9,7 @@ namespace Microsoft.Extensions.DependencyInjection
     /// Requires v8 or later of Microsoft.Extensions.DependencyInjection package.
     /// </summary>
     /// <typeparam name="TKey">Type of service key.</typeparam>
-    [AttributeUsage(AttributeTargets.Class)]
+    [AttributeUsage(AttributeTargets.Class, AllowMultiple = true)]
     partial class ServiceAttribute<TKey> : Attribute
     {
         /// <summary>