added open generic support for decorator func

Author: Kamil Zakiev

src/Scrutor/ServiceCollectionExtensions.Decoration.cs

@@ -170,6 +170,11 @@ public static IServiceCollection Decorate(this IServiceCollection services, Type
             Preconditions.NotNull(serviceType, nameof(serviceType));
             Preconditions.NotNull(decorator, nameof(decorator));
 
+            if (serviceType.IsOpenGeneric())
+            {
+                return services.DecorateOpenGeneric(serviceType, decorator);
+            }
+
             return services.DecorateDescriptors(serviceType, x => x.Decorate(decorator));
         }
 
@@ -243,10 +248,28 @@ private static bool IsSameGenericType(Type t1, Type t2)
             return t1.IsGenericType && t2.IsGenericType && t1.GetGenericTypeDefinition() == t2.GetGenericTypeDefinition();
         }
 
-        private static bool TryDecorateOpenGeneric(this IServiceCollection services, Type serviceType, Type decoratorType)
+        private static IServiceCollection DecorateOpenGeneric(this IServiceCollection services, Type serviceType, Func<object, IServiceProvider, object> decorator)
         {
             bool TryDecorate(Type[] typeArguments)
             {
+                var closedServiceType = serviceType.MakeGenericType(typeArguments);
+                return services.TryDecorateDescriptors(closedServiceType, x => x.Decorate(decorator));
+            }
+
+            if (services.TryDecorateOpenGeneric(serviceType, openTypeDecorator: TryDecorate))
+            {
+                return services;
+            }
+
+            throw new MissingTypeRegistrationException(serviceType);
+        }
+
+        private static bool TryDecorateOpenGeneric(this IServiceCollection services, Type serviceType, Type decoratorType = null, Func<Type[], bool> openTypeDecorator = null)
+        {
+            bool TryDecorate(Type[] typeArguments)
+            {
+                Preconditions.NotNull(decoratorType, nameof(decoratorType));
+
                 var closedServiceType = serviceType.MakeGenericType(typeArguments);
                 var closedDecoratorType = decoratorType.MakeGenericType(typeArguments);
 
@@ -263,7 +286,8 @@ bool TryDecorate(Type[] typeArguments)
                 return false;
             }
 
-            return arguments.Aggregate(true, (result, args) => result && TryDecorate(args));
+            var tryDecorate = openTypeDecorator ?? TryDecorate;
+            return arguments.Aggregate(true, (result, args) => result && tryDecorate(args));
         }
 
         private static IServiceCollection DecorateDescriptors(this IServiceCollection services, Type serviceType, Func<ServiceDescriptor, ServiceDescriptor> decorator)

test/Scrutor.Tests/OpenGenericDecorationTests.cs

@@ -41,6 +41,28 @@ public void CanDecorateOpenGenericTypeBasedOnInterface()
             Assert.IsType<MyQueryHandler>(loggingDecorator.Inner);
         }
 
+        [Fact]
+        public void CanDecorateOpenGenericTypeBasedOnInterfaceByDecoratorFunc()
+        {
+            var provider = ConfigureProvider(services =>
+            {
+                services.AddSingleton<IQueryHandler<MyQuery, MyResult>, MySpecialQueryHandler>();
+                services.Decorate(typeof(IQueryHandler<,>), (handlerObj, serviceProvider) =>
+                {
+                    if (handlerObj is ISpecialInterface specialInterface)
+                    {
+                        specialInterface.InitSomeField();
+                    }
+
+                    return handlerObj;
+                });
+            });
+
+            var instance = provider.GetRequiredService<IQueryHandler<MyQuery, MyResult>>();
+            var myQueryHandler = Assert.IsType<MySpecialQueryHandler>(instance);
+            Assert.True(myQueryHandler.GetSomeField());
+        }
+
         [Fact]
         public void DecoratingNonRegisteredOpenGenericServiceThrows()
         {
@@ -79,6 +101,22 @@ public void DecoratingOpenGenericTypeBasedOnGrandparentInterfaceDoesNotDecorateP
         }
     }
 
+    public interface ISpecialInterface
+    {
+        void InitSomeField();
+    }
+
+    public class MySpecialQueryHandler : QueryHandler<MyQuery, MyResult>, ISpecialInterface
+    {
+        private bool _someField = false;
+        public void InitSomeField()
+        {
+            _someField = true;
+        }
+
+        public bool GetSomeField() => _someField;
+    }
+
     public class MyQuery { }
 
     public class MyResult { }