diff --git a/PearlSourceGenerators/PearlSourceGenerators/AutoInjectionSourceGenerator.cs b/PearlSourceGenerators/PearlSourceGenerators/AutoInjectionSourceGenerator.cs
index 742c500..616aa0c 100644
--- a/PearlSourceGenerators/PearlSourceGenerators/AutoInjectionSourceGenerator.cs
+++ b/PearlSourceGenerators/PearlSourceGenerators/AutoInjectionSourceGenerator.cs
@@ -7,6 +7,7 @@
 using Microsoft.CodeAnalysis.CSharp;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Text;
+using PearlSourceGenerators.Utility;
 
 namespace PearlSourceGenerators;
 
@@ -52,42 +53,13 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
     private static (ClassDeclarationSyntax, bool attributesFound) GetClassDeclarationForSourceGen(
         GeneratorSyntaxContext context)
     {
-        var hasClassAttr = false;
-        var hasPropAttr = false;
-        
         // check if the class has the attribute we expect
         var classDeclarationSyntax = (ClassDeclarationSyntax)context.Node;
-        foreach (var attributeListSyntax in classDeclarationSyntax.AttributeLists)
-        foreach (var attributeSyntax in attributeListSyntax.Attributes)
-        {
-            if (ModelExtensions.GetSymbolInfo(context.SemanticModel, attributeSyntax).Symbol is not IMethodSymbol attributeSymbol)
-                continue; // ignore symbols we can't get
-            var attrName = attributeSymbol.ContainingType.ToDisplayString();
-            
-            if (attrName != $"{Namespace}.{ClassAttribute}") 
-                continue;
-            
-            hasClassAttr = true;
-            break;
-        }
+        var hasClassAttr = classDeclarationSyntax.HasAttribute(context, $"{Namespace}.{ClassAttribute}");
         
-        // check if any fields have the attribute we expect
-        var fields = classDeclarationSyntax.Members
-            .Where(m => m is PropertyDeclarationSyntax);
-        foreach (var field in fields)
-        foreach (var attributeListSyntax in field.AttributeLists)
-        foreach (var attributeSyntax in attributeListSyntax.Attributes)
-        {
-            if (ModelExtensions.GetSymbolInfo(context.SemanticModel, attributeSyntax).Symbol is not IMethodSymbol attributeSymbol)
-                continue;
-            var attrName = attributeSymbol.ContainingType.ToDisplayString();
-
-            if (attrName != $"{Namespace}.{PropertyAttribute}")
-                continue;
-            
-            hasPropAttr = true;
-            break;
-        }
+        // check if any properties have the attribute we expect
+        var hasPropAttr = classDeclarationSyntax.AnyClassMembersHave<PropertyDeclarationSyntax>(context,
+            $"{Namespace}.{PropertyAttribute}");
         
         return (classDeclarationSyntax, hasClassAttr && hasPropAttr);
     }
@@ -106,17 +78,7 @@ private void GenerateCode(SourceProductionContext context, Compilation compilati
             
             var namespaceName = classSymbol.ContainingNamespace.ToDisplayString();
             var className = classSymbol.Name;
-
-            if (!classDecl.Modifiers.Any(m => m.IsKind(SyntaxKind.PartialKeyword)))
-            {
-                context.ReportDiagnostic(Diagnostic.Create(new DiagnosticDescriptor(
-                    "PSG001", 
-                    "Class must be partial", 
-                    "Class must be partial", 
-                    "PearlSourceGenerators",
-                    DiagnosticSeverity.Warning, true),
-                    classDecl.GetLocation()));
-            }
+            
             var properties = classSymbol.GetMembers()
                 .OfType<IPropertySymbol>()
                 .Where(p =>
diff --git a/PearlSourceGenerators/PearlSourceGenerators/LazyDependencyInjectionSourceGenerator.cs b/PearlSourceGenerators/PearlSourceGenerators/LazyDependencyInjectionSourceGenerator.cs
index dc4d2ce..2209e0f 100644
--- a/PearlSourceGenerators/PearlSourceGenerators/LazyDependencyInjectionSourceGenerator.cs
+++ b/PearlSourceGenerators/PearlSourceGenerators/LazyDependencyInjectionSourceGenerator.cs
@@ -5,6 +5,7 @@
 using Microsoft.CodeAnalysis;
 using Microsoft.CodeAnalysis.CSharp.Syntax;
 using Microsoft.CodeAnalysis.Text;
+using PearlSourceGenerators.Utility;
 
 namespace PearlSourceGenerators;
 
@@ -55,45 +56,15 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
     private static (ClassDeclarationSyntax, bool attributesFound) GetClassDeclarationForSourceGen(
         GeneratorSyntaxContext context)
     {
-        var hasClassAttr = false;
-        var hasFieldAttr = false;
-        
         // check if the class has the attribute we expect
         var classDeclarationSyntax = (ClassDeclarationSyntax)context.Node;
-        foreach (var attributeListSyntax in classDeclarationSyntax.AttributeLists)
-        foreach (var attributeSyntax in attributeListSyntax.Attributes)
-        {
-            if (context.SemanticModel.GetSymbolInfo(attributeSyntax).Symbol is not IMethodSymbol attributeSymbol)
-                continue; // ignore symbols we can't get
-            var attrName = attributeSymbol.ContainingType.ToDisplayString();
-            
-            if (attrName != $"{Namespace}.{ClassAttribute}") 
-                continue;
-            
-            hasClassAttr = true;
-            break;
-        }
+        var hasClassAttr = classDeclarationSyntax.HasAttribute(context, $"{Namespace}.{ClassAttribute}");
 
         // check if any fields have the attribute we expect
-        var fields = classDeclarationSyntax.Members
-            .Where(m => m is FieldDeclarationSyntax);
-        foreach (var field in fields)
-        foreach (var attributeListSyntax in field.AttributeLists)
-        foreach (var attributeSyntax in attributeListSyntax.Attributes)
-        {
-            if (context.SemanticModel.GetSymbolInfo(attributeSyntax).Symbol is not IMethodSymbol attributeSymbol)
-                continue;
-            var attrName = attributeSymbol.ContainingType.ToDisplayString();
-
-            if (attrName != $"{Namespace}.{FieldAttribute}")
-                continue;
-            
-            hasFieldAttr = true;
-            break;
-        }
+        var hasFieldAttr = classDeclarationSyntax.AnyClassMembersHave<FieldDeclarationSyntax>(context,
+                $"{Namespace}.{FieldAttribute}");
         
         return (classDeclarationSyntax, hasClassAttr && hasFieldAttr);
-
     }
 
     private void GenerateCode(SourceProductionContext context, Compilation compilation,
diff --git a/PearlSourceGenerators/PearlSourceGenerators/Utility/ClassValidationExtensions.cs b/PearlSourceGenerators/PearlSourceGenerators/Utility/ClassValidationExtensions.cs
new file mode 100644
index 0000000..bb66d65
--- /dev/null
+++ b/PearlSourceGenerators/PearlSourceGenerators/Utility/ClassValidationExtensions.cs
@@ -0,0 +1,69 @@
+using System.Linq;
+using Microsoft.CodeAnalysis;
+using Microsoft.CodeAnalysis.CSharp;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
+
+namespace PearlSourceGenerators.Utility;
+
+/// <summary>
+/// Utilities for validating that a class should be processed.
+/// </summary>
+public static class GeneratorValidationExtensions
+{
+    /// <summary>
+    /// Check if a given member (ie. class, field, method, etc) has at least one instance of an attribute present.
+    /// </summary>
+    /// <param name="member">The member to check</param>
+    /// <param name="context">The syntax context of the generator</param>
+    /// <param name="attributeType">The fully qualified type name of the attribute</param>
+    /// <returns>true if the attribute is present, false if it is not</returns>
+    public static bool HasAttribute(this MemberDeclarationSyntax member, GeneratorSyntaxContext context,
+        string attributeType)
+    {
+        var attrs = member.AttributeLists
+            .SelectMany(al => al.Attributes);
+        
+        foreach (var attr in attrs)
+        {
+            if (context.SemanticModel.GetSymbolInfo(attr).Symbol is not IMethodSymbol attrSym)
+                continue;
+
+            if (attrSym.ContainingType.ToDisplayString() != attributeType)
+                continue;
+            
+            return true;
+        }
+
+        return false;
+    }
+
+    /// <summary>
+    /// Check if any class members have at least one instance of an attribute present
+    /// </summary>
+    /// <param name="cls">Class to check</param>
+    /// <param name="context">The syntax context of the source generator</param>
+    /// <param name="attributeType">The fully-qualified name of the attribute</param>
+    /// <returns>true if the attribute is present, false if it is not</returns>
+    public static bool AnyClassMembersHave(this ClassDeclarationSyntax cls, GeneratorSyntaxContext context,
+        string attributeType)
+    {
+        return cls.Members
+            .Any(p => p.HasAttribute(context, attributeType));   
+    }
+    
+    /// <summary>
+    /// Check if any class members of the given type have at least one instance of an attribute present
+    /// </summary>
+    /// <typeparam name="T">The parameter type to check</typeparam>
+    /// <param name="cls">Class to check</param>
+    /// <param name="context">The syntax context of the source generator</param>
+    /// <param name="attributeType">The fully-qualified name of the attribute</param>
+    /// <returns>true if the attribute is present, false if it is not</returns>
+    public static bool AnyClassMembersHave<T>(this ClassDeclarationSyntax cls, GeneratorSyntaxContext context,
+        string attributeType) where T : MemberDeclarationSyntax
+    {
+        return cls.Members
+            .OfType<T>()
+            .Any(p => p.HasAttribute(context, attributeType));   
+    }
+}
\ No newline at end of file