Extended ILifetimeSelector with a new WithLifetime method that accepts a Func<Type, ServiceLifetime> parameter, i.e. a function that maps the implementation type to the service lifetime

Author: ripebananas

src/Scrutor/ILifetimeSelector.cs

@@ -1,4 +1,5 @@
-using Microsoft.Extensions.DependencyInjection;
+using System;
+using Microsoft.Extensions.DependencyInjection;
 
 namespace Scrutor;
 
@@ -23,4 +24,9 @@ public interface ILifetimeSelector : IServiceTypeSelector
     /// Registers each matching concrete type with the specified <paramref name="lifetime"/>.
     /// </summary>
     IImplementationTypeSelector WithLifetime(ServiceLifetime lifetime);
+
+    /// <summary>
+    /// Registers each matching concrete type with the specified <paramref name="lifetime"/>.
+    /// </summary>
+    IImplementationTypeSelector WithLifetime(Func<Type, ServiceLifetime> lifetime);
 }

src/Scrutor/LifetimeSelector.cs

@@ -21,7 +21,7 @@ public LifetimeSelector(ServiceTypeSelector inner, IEnumerable<TypeMap> typeMaps
 
     private IEnumerable<TypeFactoryMap> TypeFactoryMaps { get; }
 
-    public ServiceLifetime? Lifetime { get; set; }
+    public Func<Type, ServiceLifetime>? Lifetime { get; set; }
 
     public IImplementationTypeSelector WithSingletonLifetime()
     {
@@ -37,11 +37,17 @@ public IImplementationTypeSelector WithTransientLifetime()
     {
         return WithLifetime(ServiceLifetime.Transient);
     }
-
     public IImplementationTypeSelector WithLifetime(ServiceLifetime lifetime)
     {
         Preconditions.IsDefined(lifetime, nameof(lifetime));
 
+        return WithLifetime(_ => lifetime);
+    }
+
+    public IImplementationTypeSelector WithLifetime(Func<Type, ServiceLifetime> lifetime)
+    {
+        Preconditions.NotNull(lifetime, nameof(lifetime));
+
         Inner.PropagateLifetime(lifetime);
 
         return this;
@@ -205,7 +211,7 @@ void ISelector.Populate(IServiceCollection services, RegistrationStrategy? strat
     {
         strategy ??= RegistrationStrategy.Append;
 
-        var lifetime = Lifetime ?? ServiceLifetime.Transient;
+        var serviceLifetimes = new Dictionary<Type, ServiceLifetime>();
 
         foreach (var typeMap in TypeMaps)
         {
@@ -218,7 +224,7 @@ void ISelector.Populate(IServiceCollection services, RegistrationStrategy? strat
                     throw new InvalidOperationException($@"Type ""{implementationType.ToFriendlyName()}"" is not assignable to ""${serviceType.ToFriendlyName()}"".");
                 }
 
-                var descriptor = new ServiceDescriptor(serviceType, implementationType, lifetime);
+                var descriptor = new ServiceDescriptor(serviceType, implementationType, GetOrAddLifetime(serviceLifetimes, implementationType));
 
                 strategy.Apply(services, descriptor);
             }
@@ -228,10 +234,17 @@ void ISelector.Populate(IServiceCollection services, RegistrationStrategy? strat
         {
             foreach (var serviceType in typeFactoryMap.ServiceTypes)
             {
-                var descriptor = new ServiceDescriptor(serviceType, typeFactoryMap.ImplementationFactory, lifetime);
+                var descriptor = new ServiceDescriptor(serviceType, typeFactoryMap.ImplementationFactory, GetOrAddLifetime(serviceLifetimes, typeFactoryMap.ImplementationType));
 
                 strategy.Apply(services, descriptor);
             }
         }
     }
+
+    private ServiceLifetime GetOrAddLifetime(Dictionary<Type, ServiceLifetime> serviceLifetimes, Type implementationType)
+    {
+        return serviceLifetimes.TryGetValue(implementationType, out var lifetime)
+            ? lifetime
+            : (serviceLifetimes[implementationType] = Lifetime?.Invoke(implementationType) ?? ServiceLifetime.Transient);
+    }
 }

src/Scrutor/ServiceTypeSelector.cs

@@ -71,7 +71,7 @@ public ILifetimeSelector AsSelfWithInterfaces(Func<Type, bool> predicate)
 
         return AddSelector(
             Types.Select(t => new TypeMap(t, new[] { t })),
-            Types.Select(t => new TypeFactoryMap(x => x.GetRequiredService(t), Selector(t, predicate))));
+            Types.Select(t => new TypeFactoryMap(x => x.GetRequiredService(t), Selector(t, predicate), t)));
 
         static IEnumerable<Type> Selector(Type type, Func<Type, bool> predicate)
         {
@@ -209,7 +209,7 @@ public IServiceTypeSelector AddClasses(Action<IImplementationTypeFilter> action,
 
     #endregion
 
-    internal void PropagateLifetime(ServiceLifetime lifetime)
+    internal void PropagateLifetime(Func<Type, ServiceLifetime> lifetime)
     {
         foreach (var selector in Selectors.OfType<LifetimeSelector>())
         {

src/Scrutor/TypeFactoryMap.cs

@@ -5,13 +5,17 @@ namespace Scrutor;
 
 internal struct TypeFactoryMap
 {
-    public TypeFactoryMap(Func<IServiceProvider, object> implementationFactory, IEnumerable<Type> serviceTypes)
+    public TypeFactoryMap(Func<IServiceProvider, object> implementationFactory, IEnumerable<Type> serviceTypes, Type implementationType)
     {
         ImplementationFactory = implementationFactory;
         ServiceTypes = serviceTypes;
+        ImplementationType = implementationType;
     }
 
     public Func<IServiceProvider, object> ImplementationFactory { get; }
 
     public IEnumerable<Type> ServiceTypes { get; }
+
+    public Type ImplementationType { get; }
+
 }