From b14d836fcfd4738e3a750c7d22b4af5556516ea6 Mon Sep 17 00:00:00 2001
From: NikolayPianikov <Nikolay.Pianikov@jetbrains.com>
Date: Fri, 21 Jan 2022 13:09:43 +0300
Subject: [PATCH] Add ExcludeFromCodeCoverage to extensions

---
 Immutype.Tests/Integration/TestExtensions.cs |  1 +
 Immutype/Core/ExtensionsFactory.cs           | 14 ++++++++++++++
 2 files changed, 15 insertions(+)

diff --git a/Immutype.Tests/Integration/TestExtensions.cs b/Immutype.Tests/Integration/TestExtensions.cs
index 8e5b838..4687f0a 100644
--- a/Immutype.Tests/Integration/TestExtensions.cs
+++ b/Immutype.Tests/Integration/TestExtensions.cs
@@ -13,6 +13,7 @@
 namespace Immutype.Tests.Integration
 {
     using System.Collections.Immutable;
+    using System.Diagnostics.CodeAnalysis;
     using System.Threading;
 
     public static class TestExtensions
diff --git a/Immutype/Core/ExtensionsFactory.cs b/Immutype/Core/ExtensionsFactory.cs
index 4307ba8..3232f7a 100644
--- a/Immutype/Core/ExtensionsFactory.cs
+++ b/Immutype/Core/ExtensionsFactory.cs
@@ -51,6 +51,8 @@ public IEnumerable<Source> Create(GenerationContext<TypeDeclarationSyntax> conte
                 .AddModifiers(SyntaxFactory.Token(SyntaxKind.StaticKeyword), SyntaxFactory.Token(SyntaxKind.PartialKeyword))
                 .AddMembers(_methodsFactory.Create(context, typeSyntax, parameters).ToArray());
 
+            extensionsClass = TryAddAttribute(context.SemanticModel, extensionsClass, "System.Diagnostics.CodeAnalysis.ExcludeFromCodeCoverage");
+
             var code = CreateRootNode(typeDeclarationSyntax, AdditionalUsings, extensionsClass).NormalizeWhitespace().ToString();
             var fileName = string.Join(".", ns.Select(i => i.Name.ToString()).Concat(new []{typeDeclarationSyntax.Identifier.Text}));
             yield return new Source(fileName, SourceText.From(code, Encoding.UTF8));
@@ -82,5 +84,17 @@ private static UsingDirectiveSyntax[] GetUsings(IEnumerable<UsingDirectiveSyntax
             var currentUsins = usings.Select(i => i.Name.ToString()).ToImmutableHashSet();
             return additionalUsings.Where(i => !currentUsins.Contains(i.Name.ToString())).ToArray();
         }
+        
+        private static ClassDeclarationSyntax TryAddAttribute(SemanticModel semanticModel, ClassDeclarationSyntax classDeclarationSyntax, string attributeClassName)
+        {
+            var excludeFromCodeCoverageType = semanticModel.Compilation.GetTypeByMetadataName(attributeClassName + "Attribute");
+            if (excludeFromCodeCoverageType != default)
+            {
+                classDeclarationSyntax = classDeclarationSyntax
+                    .AddAttributeLists(SyntaxFactory.AttributeList().AddAttributes(SyntaxFactory.Attribute(SyntaxFactory.IdentifierName(attributeClassName))));
+            }
+
+            return classDeclarationSyntax;
+        }
     }
 }
\ No newline at end of file