Exclude classes from assembly-level discovery

Author: ashmind

Generators/Internal/KnownTypes.cs

@@ -15,6 +15,15 @@ public static class NamedParameters {
             }
         }
 
+        public static class GenerateMocksForTypesAttribute {
+            public const string Name = "GenerateMocksForTypesAttribute";
+            public const string NameWithoutAttribute = "GenerateMocksForTypes";
+
+            public static bool NamespaceMatches(INamespaceSymbol? @namespace) => @namespace is {
+                Name: "SourceMock", ContainingNamespace: { IsGlobalNamespace: true }
+            };
+        }
+
         public static class IMock {
             public const string FullName = "SourceMock.IMock";
         }

Generators/Internal/MockTargetDiscovery.cs

@@ -17,7 +17,7 @@ public IEnumerable<MockTargetMember> GetMembersToMock(MockTarget target) {
             #pragma warning restore HAA0502
 
             var memberId = 1;
-            foreach (var member in target.PotentiallyLoadedMembers ?? target.Type.GetMembers()) {
+            foreach (var member in target.Type.GetMembers()) {
                 seen.Add(member.Name);
 
                 if (GetTargetMember(member, memberId) is not {} discovered)

Generators/Internal/Models/MockTarget.cs

@@ -1,4 +1,3 @@
-using System.Collections.Immutable;
 using Microsoft.CodeAnalysis;
 using Roslyn.Utilities;
 
@@ -7,17 +6,14 @@ internal readonly struct MockTarget {
         [PerformanceSensitive("")]
         public MockTarget(
             INamedTypeSymbol targetType,
-            string targetTypeQualifiedName,
-            ImmutableArray<ISymbol>? potentiallyLoadedMembers
+            string targetTypeQualifiedName
         )
         {
             Type = targetType;
             FullTypeName = targetTypeQualifiedName;
-            PotentiallyLoadedMembers = potentiallyLoadedMembers;
         }
 
         public INamedTypeSymbol Type { get; }
         public string FullTypeName { get; }
-        public ImmutableArray<ISymbol>? PotentiallyLoadedMembers { get; }
     }
 }

Generators/MockGenerator.cs

@@ -15,11 +15,11 @@ internal class MockGenerator : ISourceGenerator, IDisposable {
         private static class DiagnosticDescriptors {
             #pragma warning disable RS2008 // Enable analyzer release tracking
             public static readonly DiagnosticDescriptor SingleMockFailedToGenerate = new(
-                "SM0001", "Failed to generate a single mock", "Failed to generate mock for {0}: {1}",
+                "SourceMock001", "Failed to generate a single mock", "Failed to generate mock for {0}: {1}",
                 "Generation", DiagnosticSeverity.Warning, isEnabledByDefault: true
             );
             public static readonly DiagnosticDescriptor RegexPatternFailedToParse = new(
-                "SM0011", "Regex pattern is not a valid regex", "Regex pattern \"{0}\" cannot be parsed: {1}",
+                "SourceMock002", "Regex pattern is not a valid regex", "Regex pattern \"{0}\" cannot be parsed: {1}",
                 "Generation", DiagnosticSeverity.Error, isEnabledByDefault: true
             );
             #pragma warning restore RS2008 // Enable analyzer release tracking
@@ -44,13 +44,7 @@ public void Execute(GeneratorExecutionContext context) {
                 var attributes = context.Compilation.Assembly.GetAttributes();
                 GeneratorLog.Log("Get attributes finished");
                 foreach (var attribute in attributes) {
-                    if (attribute.AttributeClass is not { Name: KnownTypes.GenerateMocksForAssemblyOfAttribute.Name } attributeClass)
-                        continue;
-
-                    if (!KnownTypes.GenerateMocksForAssemblyOfAttribute.NamespaceMatches(attributeClass.ContainingNamespace))
-                        continue;
-
-                    GenerateMocksForAttributeTargetAssembly(attribute, context);
+                    ProcessAssemblyAttribute(attribute, context);
                 }
             }
             catch (Exception ex) {
@@ -60,6 +54,33 @@ public void Execute(GeneratorExecutionContext context) {
             GeneratorLog.Log("MockGenerator.Execute finished");
         }
 
+        private void ProcessAssemblyAttribute(AttributeData attribute, in GeneratorExecutionContext context) {
+            if (attribute.AttributeClass is not {} attributeClass)
+                return;
+
+            switch (attributeClass.Name) {
+                case KnownTypes.GenerateMocksForAssemblyOfAttribute.Name:
+                    if (!KnownTypes.GenerateMocksForAssemblyOfAttribute.NamespaceMatches(attributeClass.ContainingNamespace))
+                        return;
+                    GenerateMocksForAttributeTargetAssembly(attribute, context);
+                    break;
+
+                case KnownTypes.GenerateMocksForTypesAttribute.Name:
+                    if (!KnownTypes.GenerateMocksForTypesAttribute.NamespaceMatches(attributeClass.ContainingNamespace))
+                        return;
+
+                    if (attribute.ConstructorArguments.ElementAtOrDefault(0) is not { Kind: TypedConstantKind.Array, Values: var typeofConstants })
+                        return;
+
+                    foreach (var typeofConstant in typeofConstants) {
+                        if (typeofConstant is not { Value: INamedTypeSymbol type })
+                            continue;
+                        GenerateMockForType(new MockTarget(type, GetFullTypeName(type)), assemblyCacheBuilder: null, context);
+                    }
+                    break;
+            }
+        }
+
         [PerformanceSensitive("")]
         private void GenerateMocksForAttributeTargetAssembly(AttributeData attribute, in GeneratorExecutionContext context) {
             // intermediate code state? just in case
@@ -118,7 +139,9 @@ in GeneratorExecutionContext context
             #pragma warning restore HAA0401
                 switch (member) {
                     case INamedTypeSymbol type:
-                        GenerateMockForTypeIfApplicable(type, excludeRegex, assemblyCacheBuilder, context);
+                        if (!ShouldIncludeInMocksForAssembly(type, excludeRegex, out var fullName, context))
+                            continue;
+                        GenerateMockForType(new MockTarget(type, fullName!), assemblyCacheBuilder, context);
                         break;
 
                     case INamespaceSymbol nested:
@@ -129,59 +152,45 @@ in GeneratorExecutionContext context
         }
 
         [PerformanceSensitive("")]
-        private void GenerateMockForTypeIfApplicable(
+        private bool ShouldIncludeInMocksForAssembly(
             INamedTypeSymbol type,
             Regex? excludeRegex,
-            ImmutableArray<(string, SourceText)>.Builder assemblyCacheBuilder,
+            out string? fullName,
             in GeneratorExecutionContext context
         ) {
-            if (!ShouldGenerateMockForTypeKind(type, out var potentiallyLoadedMembers))
-                return;
+            fullName = null;
+            if (type.TypeKind != TypeKind.Interface)
+                return false;
 
             if (type.DeclaredAccessibility != Accessibility.Public) {
                 var isVisibleInternal = type.DeclaredAccessibility == Accessibility.Internal
                                      && type.ContainingAssembly.GivesAccessTo(context.Compilation.Assembly);
                 if (!isVisibleInternal)
-                    return;
+                    return false;
             }
 
-            var fullName = type.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat);
+            fullName = GetFullTypeName(type);
             if (excludeRegex != null && excludeRegex.IsMatch(fullName))
-                return;
+                return false;
 
-            GenerateMockForType(new MockTarget(type, fullName, potentiallyLoadedMembers), assemblyCacheBuilder, context);
+            return true;
         }
 
         [PerformanceSensitive("")]
-        private bool ShouldGenerateMockForTypeKind(INamedTypeSymbol type, out ImmutableArray<ISymbol>? members) {
-            members = null;
-            if (type.TypeKind == TypeKind.Interface)
-                return true;
-
-            if (type.TypeKind == TypeKind.Class) {
-                if (type.IsSealed)
-                    return false;
-
-                members = type.GetMembers();
-                foreach (var member in members) {
-                    if (member.IsVirtual || member.IsAbstract)
-                        return true;
-                }
-            }
-
-            return false;
+        private static string GetFullTypeName(INamedTypeSymbol type) {
+            return type.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat);
         }
 
         [PerformanceSensitive("")]
         private void GenerateMockForType(
             MockTarget target,
-            ImmutableArray<(string, SourceText)>.Builder assemblyCacheBuilder,
+            ImmutableArray<(string, SourceText)>.Builder? assemblyCacheBuilder,
             in GeneratorExecutionContext context
         ) {
             if (_mockedTypeCache.TryGetValue(target.Type, out var cached)) {
                 GeneratorLog.Log("Using cached mock for type " + target.FullTypeName);
                 context.AddSource(cached.name, cached.source);
-                assemblyCacheBuilder.Add(cached);
+                assemblyCacheBuilder?.Add(cached);
                 return;
             }
 
@@ -204,7 +213,7 @@ in GeneratorExecutionContext context
             var mockSource = SourceText.From(mockContent, Encoding.UTF8);
             context.AddSource(mockFileName, mockSource);
             _mockedTypeCache.TryAdd(target.Type, (mockFileName, mockSource));
-            assemblyCacheBuilder.Add((mockFileName, mockSource));
+            assemblyCacheBuilder?.Add((mockFileName, mockSource));
         }
 
         public void Dispose() {

Tests/MockManifest.cs

@@ -1,4 +1,5 @@
 using SourceMock;
 using SourceMock.Tests.Interfaces;
 
-[assembly: GenerateMocksForAssemblyOf(typeof(IMockable), ExcludeRegex = "ExcludedInterface")]
+[assembly: GenerateMocksForAssemblyOf(typeof(IMockable), ExcludeRegex = "ExcludedInterface")]
+[assembly: GenerateMocksForTypes(typeof(AbstractClass))]

[Main]/GenerateMocksForAssemblyOfAttribute.cs

@@ -2,14 +2,17 @@
 using System.Text.RegularExpressions;
 
 namespace SourceMock {
-    /// <summary>Requests mock generation for a specific assembly.</summary>
+    /// <summary>
+    /// Requests mock generation for all interfaces in the specified assembly.
+    /// </summary>
     /// <remarks>
-    /// When this attribute is used in a test assembly, SourceMock will generate
-    /// mocks for all interfaces in the target assembly and include those into
-    /// the test assembly.
+    /// When this attribute is used in a test project, SourceMock will generate
+    /// mocks for all interfaces in the target assembly, and include those into
+    /// the test project.
     /// </remarks>
     [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
-    public class GenerateMocksForAssemblyOfAttribute : Attribute {
+    public class GenerateMocksForAssemblyOfAttribute
+        : Attribute {
         /// <summary>
         /// Initializes a new instance of the <see cref="GenerateMocksForAssemblyOfAttribute" /> class.
         /// </summary>

[Main]/GenerateMocksForTypesAttribute.cs

@@ -0,0 +1,19 @@
+using System;
+
+namespace SourceMock {
+    /// <summary>
+    /// Requests mock generation for specified classes or interfaces.
+    /// </summary>
+    [AttributeUsage(AttributeTargets.Assembly, AllowMultiple = true)]
+    public class GenerateMocksForTypesAttribute : Attribute {
+        /// <summary>
+        /// Initializes a new instance of the <see cref="GenerateMocksForTypesAttribute" /> class.
+        /// </summary>
+        /// <param name="types">Types to generate mocks for.</param>
+        public GenerateMocksForTypesAttribute(params Type[] types) {
+            Types = types;
+        }
+
+        internal Type[] Types { get; }
+    }
+}