diff --git a/src/BlueWest.MapTo/EfAddMethodsGenerator.cs b/src/BlueWest.MapTo/EfAddMethodsGenerator.cs
index 96ad530..bc2f61f 100644
--- a/src/BlueWest.MapTo/EfAddMethodsGenerator.cs
+++ b/src/BlueWest.MapTo/EfAddMethodsGenerator.cs
@@ -10,15 +10,6 @@ using Microsoft.CodeAnalysis.CSharp.Syntax;
 
 namespace MapTo
 {
-    /*/// <summary>
-    /// Base source code generator typed class.
-    /// </summary>
-    public class SourceCodeGenerator : ISourceGenerator
-    {
-        public virtual void Initialize(GeneratorInitializationContext context) {}
-        public virtual void Execute(GeneratorExecutionContext context) {}
-    }*/
-    
     /// <summary>
     /// Ef methods source generator
     /// </summary>
@@ -72,7 +63,7 @@ namespace MapTo
                         .ToString() ?? string.Empty;
 
 
-                var mappingContext = EfAddGeneratorContext.Create(compilation, options, candidateMember.MemberDeclarationSyntax);
+                var mappingContext = EfGeneratorContext.Create(compilation, options, candidateMember.MemberDeclarationSyntax);
                 
                 mappingContext.Diagnostics.ForEach(context.ReportDiagnostic);
                 
diff --git a/src/BlueWest.MapTo/EfAddGeneratorContext.cs b/src/BlueWest.MapTo/EfGeneratorContext.cs
similarity index 86%
rename from src/BlueWest.MapTo/EfAddGeneratorContext.cs
rename to src/BlueWest.MapTo/EfGeneratorContext.cs
index 5150473..5d5f823 100644
--- a/src/BlueWest.MapTo/EfAddGeneratorContext.cs
+++ b/src/BlueWest.MapTo/EfGeneratorContext.cs
@@ -15,11 +15,11 @@ using CSharpExtensions = Microsoft.CodeAnalysis.CSharpExtensions;
 
 namespace MapTo
 {
-    internal class EfAddGeneratorContext
+    internal class EfGeneratorContext
     {
         private readonly List<SymbolDisplayPart> _ignoredNamespaces;
 
-        protected EfAddGeneratorContext(
+        protected EfGeneratorContext(
             Compilation compilation, 
             SourceGenerationOptions sourceGenerationOptions, 
             MemberDeclarationSyntax memberSyntax)
@@ -51,14 +51,14 @@ namespace MapTo
 
         protected ImmutableArray<string> Usings { get; private set; }
 
-        public static EfAddGeneratorContext Create(
+        public static EfGeneratorContext Create(
             Compilation compilation, 
             SourceGenerationOptions sourceGenerationOptions, 
             MemberDeclarationSyntax typeSyntax)
         {
-            EfAddGeneratorContext context = typeSyntax switch
+            EfGeneratorContext context = typeSyntax switch
             {
-                PropertyDeclarationSyntax => new EfAddGeneratorContext(compilation, sourceGenerationOptions, typeSyntax),
+                PropertyDeclarationSyntax => new EfGeneratorContext(compilation, sourceGenerationOptions, typeSyntax),
                 _ => throw new ArgumentOutOfRangeException()
             };
 
@@ -100,10 +100,42 @@ namespace MapTo
             return sourceSymbol;
         }
 
+        protected TypeInfo FindSecondTypeInfoOfMemberExpressionSyntax(SemanticModel semanticModel, SyntaxNode attributeSyntax)
+        {
+            semanticModel ??= Compilation.GetSemanticModel(attributeSyntax.SyntaxTree);
+            var descendentNodes = attributeSyntax
+                .DescendantNodes();
+
+            var expression = descendentNodes.OfType<MemberAccessExpressionSyntax>().FirstOrDefault()?.Expression;
+            if (expression != null)
+            {
+                return semanticModel.GetTypeInfo(expression);
+            }
+
+            return new TypeInfo();
+        }
+        
+        protected TypeInfo? FindFirstTypeInfoOfMemberExpressionSyntax(SemanticModel semanticModel, SyntaxNode attributeSyntax)
+        {
+            semanticModel ??= Compilation.GetSemanticModel(attributeSyntax.SyntaxTree);
+            var descendentNodes = attributeSyntax
+                .DescendantNodes();
+
+            // Get the type of the key member id
+            var expression = descendentNodes.OfType<MemberAccessExpressionSyntax>().FirstOrDefault();
+            if (expression != null)
+            {
+                return semanticModel.GetTypeInfo(expression);
+            }
+
+            return null;
+        }
+        
+        
         protected string ExtractNameOfMemberName(MemberDeclarationSyntax memberDeclarationSyntax, string attributeName, SemanticModel? semanticModel = null)
         {
             var attributeData = memberDeclarationSyntax.GetAttribute(attributeName);
-            var sourceSymbol = ExtractNameOfMemberName(attributeData);
+            var sourceSymbol = ExtractNameOfMemberName(attributeData, semanticModel);
             return sourceSymbol;
         }
 
@@ -154,7 +186,7 @@ namespace MapTo
             return sourceTypeExpressionSyntax;
         }
 
-        protected string ExtractNameOfMemberName(SyntaxNode? attributeSyntax)
+        protected string ExtractNameOfMemberName(SyntaxNode? attributeSyntax, SemanticModel semanticModel)
         {
             if (attributeSyntax is null)
             {
@@ -306,7 +338,16 @@ namespace MapTo
                 GetEntityTypeSymbol(MemberSyntax, EfUpdateMethodsAttributeSource.AttributeName, semanticModel);
 
             var keyPropertyName = ExtractNameOfMemberName(MemberSyntax, EfUpdateMethodsAttributeSource.AttributeName, semanticModel);
+            
+            var keyMemberType = FindFirstTypeInfoOfMemberExpressionSyntax(semanticModel, MemberSyntax);
 
+            var secondTypeInfo = FindSecondTypeInfoOfMemberExpressionSyntax(semanticModel, MemberSyntax);
+
+            
+            
+            var keyMemberFullName = keyMemberType.Value.Type.ToDisplayString();
+            var keyMemberName = keyMemberType.Value.Type.Name;
+            
             // Try grabbing string literal if there's no nameof in it
             if (keyPropertyName == string.Empty)
             {
@@ -321,29 +362,24 @@ namespace MapTo
             string updateTypeFullName = entityTypeFullName;
             string returnTypeIdentifierName = entityTypeName;
             string returnTypeFullName = entityTypeFullName;
-            string keyPropertyTypeFullName = "int";
             
             if (efAddAttributeTypeSymbols.Length > 0)
             {
                 updateTypeIdentifierName = efAddAttributeTypeSymbols[0].Name;
                 updateTypeFullName = efAddAttributeTypeSymbols[0].ToDisplayString();
             }
-            if (efAddAttributeTypeSymbols.Length > 0)
+            
+            // Grab return type from attribute argument
+            if (efAddAttributeTypeSymbols.Length > 1)
             {
                 returnTypeIdentifierName = efAddAttributeTypeSymbols[1].Name;
                 returnTypeFullName = efAddAttributeTypeSymbols[1].ToDisplayString();
             }
-     
-
-
-            if (efAddAttributeTypeSymbols.Length > 2)
-            {
-                keyPropertyTypeFullName = efAddAttributeTypeSymbols[2].ToDisplayString();
-            }
+            
 
 
             return new EfUpdateMethodsModel(updateTypeFullName, updateTypeIdentifierName, returnTypeFullName,
-                returnTypeIdentifierName, keyPropertyName, keyPropertyTypeFullName);
+                returnTypeIdentifierName, keyPropertyName, keyMemberFullName);
         }
         
         
diff --git a/src/BlueWest.MapTo/Sources/EfMethodsSource.cs b/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
index 91d2661..8b1442d 100644
--- a/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
+++ b/src/BlueWest.MapTo/Sources/EfMethodsSource.cs
@@ -1,3 +1,4 @@
+using System;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -60,6 +61,19 @@ namespace MapTo.Sources
             return new(generatedCode, hintName);
         }
 
+        private static SourceBuilder ParseTemplate(this SourceBuilder builder, string finalTemplate)
+        {
+            string[] array = finalTemplate.Split(Environment.NewLine.ToCharArray(), StringSplitOptions.RemoveEmptyEntries);
+            foreach (var line in array)
+            { 
+
+                builder.WriteLine(line);
+            }
+
+
+            return builder;
+        }
+
         private static SourceBuilder EfAddAddEntityMethod(this SourceBuilder builder, EfMethodsModel model, string addSourceTemplate, string entityTypeName,
             string contextFullName, string entityTypeFullName, string propertyName)
         {
@@ -82,7 +96,7 @@ namespace MapTo.Sources
                     .Replace("{propertyName}", propertyName);
 
                 builder
-                    .Write(templateToSourceBuilder.ToString());
+                    .ParseTemplate(templateToSourceBuilder.ToString());
             }
 
             if (addSourceTemplate.IsEmpty())
@@ -131,7 +145,7 @@ namespace MapTo.Sources
                     .Replace("{keyVarName}", keyVarName)
                     .Replace("{existingEntityVarName}", existingVarName);
                 builder
-                    .Write(templateToSourceBuilder.ToString());
+                    .ParseTemplate(templateToSourceBuilder.ToString());
             }
 
             if (updateSourceTemplate.IsEmpty())
diff --git a/src/BlueWest.MapTo/Sources/EfUpdateMethodsAttributeSource.cs b/src/BlueWest.MapTo/Sources/EfUpdateMethodsAttributeSource.cs
index 1468c7f..0167318 100644
--- a/src/BlueWest.MapTo/Sources/EfUpdateMethodsAttributeSource.cs
+++ b/src/BlueWest.MapTo/Sources/EfUpdateMethodsAttributeSource.cs
@@ -31,7 +31,7 @@ namespace MapTo.Sources
                 .WriteOpeningBracket();
             
             builder
-                .WriteLine($"public {AttributeName}Attribute(Type updateType, Type returnType, string keyPropertyMemberName, Type keyPropertyMemberType)")
+                .WriteLine($"public {AttributeName}Attribute(Type updateType, Type returnType, string keyPropertyMemberName)")
                 .WriteOpeningBracket()
                 .WriteClosingBracket()
                 .WriteLine();