diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/ExposeAttribute.cs b/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/ExposeAttribute.cs
index 796324a..36b153d 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/ExposeAttribute.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/ExposeAttribute.cs
@@ -2,7 +2,7 @@
 
 namespace Appegy.Union
 {
-    [AttributeUsage(AttributeTargets.Struct | AttributeTargets.Class, AllowMultiple = false, Inherited = false)]
+    [AttributeUsage(AttributeTargets.Struct | AttributeTargets.Class, Inherited = false)]
     public class ExposeAttribute : Attribute
     {
         public Type[] Interfaces { get; }
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/UnionAttribute.cs b/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/UnionAttribute.cs
index 50c1eba..5257351 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/UnionAttribute.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator.Shapes/Shapes/Attributes/UnionAttribute.cs
@@ -2,7 +2,7 @@
 
 namespace Appegy.Union
 {
-    [AttributeUsage(System.AttributeTargets.Struct | AttributeTargets.Class, AllowMultiple = false, Inherited = false)]
+    [AttributeUsage(System.AttributeTargets.Struct | AttributeTargets.Class, Inherited = false)]
     public class UnionAttribute : System.Attribute
     {
         public System.Type[] Types { get; }
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator.Tests/Shapes/ShapePropertyTests.cs b/Appegy.Union.Generator~/Appegy.Union.Generator.Tests/Shapes/ShapePropertyTests.cs
index d89a5dc..4f8b98c 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator.Tests/Shapes/ShapePropertyTests.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator.Tests/Shapes/ShapePropertyTests.cs
@@ -36,7 +36,7 @@ public void WhenGettingCircle_AndTypeIsNotCircle_ThenThrowsException()
     public void WhenSettingCircle_ThenTypeIsCircle()
     {
         // Arrange
-        var shape = new Shape();
+        var shape = new Shape(new Hexagon());
         var circle = new Circle(5);
 
         // Act
@@ -76,7 +76,7 @@ public void WhenGettingRectangle_AndTypeIsNotRectangle_ThenThrowsException()
     public void WhenSettingRectangle_ThenTypeIsRectangle()
     {
         // Arrange
-        var shape = new Shape();
+        var shape = new Shape(new Hexagon());
         var rectangle = new Rectangle(4, 6);
 
         // Act
@@ -116,7 +116,7 @@ public void WhenGettingHexagon_AndTypeIsNotHexagon_ThenThrowsException()
     public void WhenSettingHexagon_ThenTypeIsHexagon()
     {
         // Arrange
-        var shape = new Shape();
+        var shape = new Shape(new Circle());
         var hexagon = new Hexagon(3);
 
         // Act
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/DiagnosticDescriptors.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/DiagnosticDescriptors.cs
index 44785ec..ec8e717 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/DiagnosticDescriptors.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/DiagnosticDescriptors.cs
@@ -6,14 +6,6 @@ public static class DiagnosticDescriptors
 {
     private const string Category = "Union";
 
-    public static DiagnosticDescriptor NotStruct { get; } = new(
-        id: "UNION001",
-        title: "Type has to be struct",
-        messageFormat: "The type '{0}' must be a struct to use [Union].",
-        category: Category,
-        defaultSeverity: DiagnosticSeverity.Error,
-        isEnabledByDefault: true);
-
     public static DiagnosticDescriptor NotPartial { get; } = new(
         id: "UNION002",
         title: "Struct is not partial",
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributeGenerator.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributeGenerator.cs
index 9693b3c..efdeb92 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributeGenerator.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributeGenerator.cs
@@ -36,11 +36,11 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             .SyntaxProvider
             .ForAttributeWithMetadataName(
                 ExposeAttributeName,
-                predicate: static (syntax, _) => syntax is StructDeclarationSyntax,
+                predicate: static (syntax, _) => syntax is TypeDeclarationSyntax,
                 transform: static (ctx, _) =>
                 {
-                    var syntax = (StructDeclarationSyntax)ctx.TargetNode;
-                    var symbol = ctx.SemanticModel.GetDeclaredSymbol(syntax) as INamedTypeSymbol;
+                    var syntax = (TypeDeclarationSyntax)ctx.TargetNode;
+                    var symbol = ctx.SemanticModel.GetDeclaredSymbol(syntax);
                     var attribute = ctx.Attributes.First();
                     var interfaces = attribute
                         .GetTypesFromConstructor(TypeKind.Interface)
@@ -53,14 +53,14 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             .SyntaxProvider
             .ForAttributeWithMetadataName(
                 UnionAttributeName,
-                predicate: static (syntax, _) => syntax is StructDeclarationSyntax,
+                predicate: static (syntax, _) => syntax is TypeDeclarationSyntax,
                 transform: static (ctx, _) =>
                 {
-                    var syntax = (StructDeclarationSyntax)ctx.TargetNode;
-                    var symbol = ctx.SemanticModel.GetDeclaredSymbol(syntax) as INamedTypeSymbol;
+                    var syntax = (TypeDeclarationSyntax)ctx.TargetNode;
+                    var symbol = ctx.SemanticModel.GetDeclaredSymbol(syntax);
                     var attribute = ctx.Attributes.First();
                     var types = attribute
-                        .GetTypesFromConstructor(TypeKind.Struct)
+                        .GetTypesFromConstructor()
                         .Select(c => new ExposeTypeInfo(c))
                         .ToImmutableList();
                     return (Symbol: symbol!, Types: types);
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributePartInput.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributePartInput.cs
index fa30373..01a46d0 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributePartInput.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/ExposeAttributePartInput.cs
@@ -4,7 +4,7 @@
 
 namespace Appegy.Union.Generator;
 
-public record struct ExposeAttributePartInput(StructDeclarationSyntax Syntax, IReadOnlyList<ExposeTypeInfo> Types, IReadOnlyList<ExposeInterfaceInfo> Interfaces);
+public record struct ExposeAttributePartInput(TypeDeclarationSyntax Syntax, IReadOnlyList<ExposeTypeInfo> Types, IReadOnlyList<ExposeInterfaceInfo> Interfaces);
 
 public readonly struct ExposeTypeInfo(INamedTypeSymbol symbol)
 {
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/Parts/ExposeDeclarationPart.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/Parts/ExposeDeclarationPart.cs
index 1291f2d..c2e4bb4 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/Parts/ExposeDeclarationPart.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Expose/Parts/ExposeDeclarationPart.cs
@@ -1,4 +1,5 @@
 using System.CodeDom.Compiler;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 
 namespace Appegy.Union.Generator;
 
@@ -8,7 +9,18 @@ public override void Generate(IndentedTextWriter codeWriter, ExposeAttributePart
     {
         var (syntax, _, interfaces) = input;
 
-        codeWriter.Write("partial struct ");
+        codeWriter.Write("partial ");
+
+        switch (syntax)
+        {
+            case StructDeclarationSyntax:
+                codeWriter.Write("struct ");
+                break;
+            case ClassDeclarationSyntax:
+                codeWriter.Write("class ");
+                break;
+        }
+
         codeWriter.Write(syntax.Identifier.Text);
 
         codeWriter.WriteLine(" :");
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/Parts/UnionDeclarationPart.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/Parts/UnionDeclarationPart.cs
index 653a00d..cc123f4 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/Parts/UnionDeclarationPart.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/Parts/UnionDeclarationPart.cs
@@ -1,4 +1,5 @@
 using System.CodeDom.Compiler;
+using Microsoft.CodeAnalysis.CSharp.Syntax;
 
 namespace Appegy.Union.Generator;
 
@@ -10,7 +11,16 @@ public override void Generate(IndentedTextWriter codeWriter, UnionAttributePartI
         var types = input.Types;
 
         codeWriter.WriteLine(AttributesNames.GeneratedCodeAttribute);
-        codeWriter.Write("partial struct ");
+        codeWriter.Write("partial ");
+        switch (syntax)
+        {
+            case StructDeclarationSyntax:
+                codeWriter.Write("struct ");
+                break;
+            case ClassDeclarationSyntax:
+                codeWriter.Write("class ");
+                break;
+        }
         codeWriter.WriteLine(syntax.Identifier.Text);
 
         codeWriter.Write("    : global::System.IEquatable<");
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeAnalyzer.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeAnalyzer.cs
index 52fe31f..e46c04f 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeAnalyzer.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeAnalyzer.cs
@@ -17,7 +17,6 @@ public class UnionAttributeAnalyzer : DiagnosticAnalyzer
             NotPartial,
             NestedNotPartial,
             NoTypesProvided,
-            NotStruct,
             DuplicateUnionType);
 
     public override void Initialize(AnalysisContext context)
@@ -39,38 +38,37 @@ private void AnalyzeAttribute(SyntaxNodeAnalysisContext context)
         VerifyPartialModifier(context, attributeSyntax);
         VerifyParentsPartial(context, attributeSyntax);
         VerifyArgumentsExistence(context, attributeSyntax);
-        VerifyAllTypesAreStruct(context, attributeSyntax);
         VerifyNoDuplicate(context, attributeSyntax);
     }
 
     private static void VerifyPartialModifier(SyntaxNodeAnalysisContext context, AttributeSyntax attributeSyntax)
     {
-        if (attributeSyntax.Parent?.Parent is not StructDeclarationSyntax structDeclaration)
+        if (attributeSyntax.Parent?.Parent is not TypeDeclarationSyntax typeDeclaration)
         {
             return;
         }
 
-        if (structDeclaration.Modifiers.Any(SyntaxKind.PartialKeyword))
+        if (typeDeclaration.Modifiers.Any(SyntaxKind.PartialKeyword))
         {
             return;
         }
 
         var diagnostic = Diagnostic.Create(
             NotPartial,
-            structDeclaration.Identifier.GetLocation(),
-            structDeclaration.Identifier.Text);
+            typeDeclaration.Identifier.GetLocation(),
+            typeDeclaration.Identifier.Text);
 
         context.ReportDiagnostic(diagnostic);
     }
 
     private static void VerifyParentsPartial(SyntaxNodeAnalysisContext context, AttributeSyntax attributeSyntax)
     {
-        if (attributeSyntax.Parent?.Parent is not StructDeclarationSyntax structDeclaration)
+        if (attributeSyntax.Parent?.Parent is not TypeDeclarationSyntax typeDeclaration)
         {
             return;
         }
 
-        var parent = structDeclaration.Parent;
+        var parent = typeDeclaration.Parent;
         while (parent is TypeDeclarationSyntax parentType)
         {
             if (!parentType.Modifiers.Any(SyntaxKind.PartialKeyword))
@@ -99,34 +97,6 @@ private void VerifyArgumentsExistence(SyntaxNodeAnalysisContext context, Attribu
         context.ReportDiagnostic(diagnostic);
     }
 
-    private void VerifyAllTypesAreStruct(SyntaxNodeAnalysisContext context, AttributeSyntax attributeSyntax)
-    {
-        var arguments = attributeSyntax.ArgumentList?.Arguments;
-        if (arguments == null || arguments.Value.Count == 0)
-        {
-            return;
-        }
-
-        foreach (var argument in arguments)
-        {
-            if (argument.Expression is not TypeOfExpressionSyntax typeOfExpression)
-            {
-                continue;
-            }
-
-            var typeInfo = context.SemanticModel.GetTypeInfo(typeOfExpression.Type);
-            if (typeInfo.Type?.TypeKind != TypeKind.Struct)
-            {
-                var diagnostic = Diagnostic.Create(
-                    NotStruct,
-                    typeOfExpression.Type.GetLocation(),
-                    typeInfo.Type?.ToDisplayString() ?? "unknown");
-
-                context.ReportDiagnostic(diagnostic);
-            }
-        }
-    }
-
     private void VerifyNoDuplicate(SyntaxNodeAnalysisContext context, AttributeSyntax attributeSyntax)
     {
         var arguments = attributeSyntax.ArgumentList?.Arguments;
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeGenerator.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeGenerator.cs
index 81e9d92..f50fc7c 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeGenerator.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributeGenerator.cs
@@ -40,14 +40,14 @@ public void Initialize(IncrementalGeneratorInitializationContext context)
             .SyntaxProvider
             .ForAttributeWithMetadataName(
                 UnionAttributeName,
-                predicate: static (syntax, _) => syntax is StructDeclarationSyntax,
+                predicate: static (syntax, _) => syntax is TypeDeclarationSyntax,
                 transform: static (ctx, _) =>
                 {
-                    var syntax = (StructDeclarationSyntax)ctx.TargetNode;
+                    var syntax = (TypeDeclarationSyntax)ctx.TargetNode;
                     var symbol = ctx.SemanticModel.GetDeclaredSymbol(syntax);
                     var attribute = ctx.Attributes.First();
                     var types = attribute
-                        .GetTypesFromConstructor(TypeKind.Struct)
+                        .GetTypesFromConstructor()
                         .Select(c => new UnionTypeInfo(c))
                         .ToImmutableList();
 
diff --git a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributePartInput.cs b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributePartInput.cs
index 83ca80b..b7f3681 100644
--- a/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributePartInput.cs
+++ b/Appegy.Union.Generator~/Appegy.Union.Generator/Union/UnionAttributePartInput.cs
@@ -4,7 +4,7 @@
 
 namespace Appegy.Union.Generator;
 
-public record struct UnionAttributePartInput(StructDeclarationSyntax Syntax, IReadOnlyList<UnionTypeInfo> Types, bool ExplicitLayout);
+public record struct UnionAttributePartInput(TypeDeclarationSyntax Syntax, IReadOnlyList<UnionTypeInfo> Types, bool ExplicitLayout);
 
 public readonly struct UnionTypeInfo(INamedTypeSymbol symbol)
 {
diff --git a/Runtime/Appegy.Union.Generator.dll b/Runtime/Appegy.Union.Generator.dll
index 3d1d07a..7ccddaf 100644
--- a/Runtime/Appegy.Union.Generator.dll
+++ b/Runtime/Appegy.Union.Generator.dll
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:8f674aa7f435e096cf702b4424fecc8de7c7a6fcdb162f2f1dd4b8fe1dcadbd2
-size 55808
+oid sha256:dfb6f14913fc4a5224b91604329dc26ddb6aceda20efd5c0f6b1d3cc263316c6
+size 55296
diff --git a/Runtime/Appegy.Union.Generator.pdb b/Runtime/Appegy.Union.Generator.pdb
index 775af18..aeedcdf 100644
Binary files a/Runtime/Appegy.Union.Generator.pdb and b/Runtime/Appegy.Union.Generator.pdb differ
diff --git a/Runtime/ExposeAttribute.cs b/Runtime/ExposeAttribute.cs
index e84d18c..36b153d 100644
--- a/Runtime/ExposeAttribute.cs
+++ b/Runtime/ExposeAttribute.cs
@@ -2,7 +2,7 @@
 
 namespace Appegy.Union
 {
-    [AttributeUsage(AttributeTargets.Struct)]
+    [AttributeUsage(AttributeTargets.Struct | AttributeTargets.Class, Inherited = false)]
     public class ExposeAttribute : Attribute
     {
         public Type[] Interfaces { get; }
diff --git a/Runtime/UnionAttribute.cs b/Runtime/UnionAttribute.cs
index e7df7b1..74bbc36 100644
--- a/Runtime/UnionAttribute.cs
+++ b/Runtime/UnionAttribute.cs
@@ -2,7 +2,7 @@
 
 namespace Appegy.Union
 {
-    [AttributeUsage(AttributeTargets.Struct)]
+    [AttributeUsage(AttributeTargets.Struct | AttributeTargets.Class, Inherited = false)]
     public class UnionAttribute : Attribute
     {
         public Type[] Types { get; }