Last active
November 13, 2018 20:34
-
-
Save VinceAvery/950360e041c4e5e0b1249e9dc8855e35 to your computer and use it in GitHub Desktop.
Add attributes to CSHarp Classes
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
public class CSharpClassRewriter : CSharpSyntaxRewriter { | |
private readonly string _sourceCode; | |
private readonly Attribute _attributeToAdd; | |
public CSharpClassRewriter( | |
string sourceCode, | |
Attribute attributeToAdd = null) { | |
if (string.IsNullOrEmpty(sourceCode)) { | |
throw new ArgumentNullException(nameof(sourceCode)); | |
} | |
_sourceCode = sourceCode; | |
_attributeToAdd = attributeToAdd; | |
} | |
public bool Rewrite(out string newSourceCode) { | |
var sourceTree = CSharpSyntaxTree.ParseText(_sourceCode); | |
var root = sourceTree.GetRoot(); | |
newSourceCode = root.ToFullString(); | |
SyntaxNode newSource = Visit(root); | |
if (newSource == root) { | |
return false; | |
} | |
newSourceCode = newSource.ToFullString(); | |
return true; | |
} | |
public override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax classNode) { | |
if (_attributeToAdd != null && !ClassHasAttribute(classNode, _attributeToAdd)) { | |
classNode = AddAttribute(classNode, _attributeToAdd); | |
} | |
return base.VisitClassDeclaration(classNode); | |
} | |
private static bool ClassHasAttribute(BaseTypeDeclarationSyntax classNode, Attribute attributeToAdd) { | |
var found = false; | |
foreach (var attributeList in classNode.AttributeLists) { | |
found = attributeList.Attributes.Any(a => a.ToFullString().Contains(attributeToAdd.GetType().FullName)); | |
} | |
return found; | |
} | |
private static ClassDeclarationSyntax AddAttribute(ClassDeclarationSyntax classDeclaration, Attribute attribute) { | |
var newAttributes = SyntaxFactory.AttributeList( | |
SyntaxFactory.SingletonSeparatedList( | |
SyntaxFactory.Attribute(SyntaxFactory.IdentifierName(attribute.GetType().FullName)) | |
) | |
); | |
classDeclaration = classDeclaration.AddAttributeLists(newAttributes); | |
return classDeclaration; | |
} | |
} | |
// Unit Tests | |
[TestClass] | |
public class CSharpClassRewriterTests { | |
/// <summary> | |
/// This tests that the required attribute is added to all classes. It also ensure that existing attributes are preserved. | |
/// </summary> | |
[TestMethod] | |
public void Class_RewriteNoRequiredAttributes_RequiredAttributeAdded() { | |
const string sourceCode = @" | |
private class TestClass1 { | |
public TestClass1() {} | |
} | |
[ExistingAttribute] | |
private class TestClass2 { | |
public TestClass2() {} | |
}"; | |
const string newSourceCodeExpected = @"[System.MTAThreadAttribute] | |
private class TestClass1 { | |
public TestClass1() {} | |
} | |
[ExistingAttribute] | |
[System.MTAThreadAttribute] private class TestClass2 { | |
public TestClass2() {} | |
}"; | |
var csharpClassRewriter = new CSharpClassRewriter(sourceCode, new MTAThreadAttribute()); | |
Assert.IsTrue(csharpClassRewriter.Rewrite(out var newSourceCode)); | |
Assert.AreEqual(newSourceCodeExpected, newSourceCode); | |
} | |
[TestMethod] | |
public void Class_RewriteAttributeExists_AttributeNotReAdded() { | |
const string sourceCode = @"[System.MTAThreadAttribute] | |
private class TestClass { | |
public TestClass() {} | |
}"; | |
var csharpClassRewriter = new CSharpClassRewriter(sourceCode, new MTAThreadAttribute()); | |
Assert.IsFalse(csharpClassRewriter.Rewrite(out _)); | |
} | |
} |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment