diff --git a/ProtoAttributor.Tests/DataContracts/DataAttributeAdderTests.cs b/ProtoAttributor.Tests/DataContracts/DataAttributeAdderTests.cs index dee0f2b..c512347 100644 --- a/ProtoAttributor.Tests/DataContracts/DataAttributeAdderTests.cs +++ b/ProtoAttributor.Tests/DataContracts/DataAttributeAdderTests.cs @@ -154,5 +154,55 @@ public void AddsAttributesWithCorrectOrderWhenFileHasProtoIgnores() output.Should().Contain("[DataMember(Order = 16)]"); _fixture.AssertOutputContainsCount(source, "[IgnoreDataMember]", 2); } + + [Fact] + public void AddsAttributesWhenFileIsEnum() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnum.cs")); + var rewriter = new DataAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("System.Runtime.Serialization"); + output.Should().Contain("[DataContract]"); + output.Should().Contain("[EnumMember]"); + _fixture.AssertOutputContainsCount(source, "[EnumMember]", 5); + + } + + [Fact] + public void AddsAttributesWhenFileIsEnumWithExistingAttributes() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithDataAttributes.cs")); + var rewriter = new DataAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("System.Runtime.Serialization"); + output.Should().Contain("[DataContract]"); + output.Should().Contain("[EnumMember]"); + _fixture.AssertOutputContainsCount(source, "[EnumMember]", 5); + } + + [Fact] + public void AddsAttributesWhenFileIsEnumWithIgnoreAttributes() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithDataMemberIgnore.cs")); + var rewriter = new DataAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("System.Runtime.Serialization"); + output.Should().Contain("[DataContract]"); + output.Should().Contain("[EnumMember]"); + _fixture.AssertOutputContainsCount(source, "[EnumMember]", 3); + _fixture.AssertOutputContainsCount(source, "[IgnoreDataMember]", 2); + } } } diff --git a/ProtoAttributor.Tests/DataContracts/DataAttributeRemoverTests.cs b/ProtoAttributor.Tests/DataContracts/DataAttributeRemoverTests.cs index be44829..894aa8c 100644 --- a/ProtoAttributor.Tests/DataContracts/DataAttributeRemoverTests.cs +++ b/ProtoAttributor.Tests/DataContracts/DataAttributeRemoverTests.cs @@ -28,6 +28,7 @@ public void AddsAttributesWithCorrectOrderWhenAttributesAlreadyExists() output.Should().NotContain("[DataContract]"); output.Should().NotContain("[KnownType"); output.Should().NotContain("[IgnoreDataMember]"); + output.Should().NotContain("[EnumMember]"); output.Should().NotContain(@"[DataMember(Order = 1, Name=""Test"")]"); output.Should().NotContain("[DataMember(Order = 2)]"); output.Should().NotContain(@"DataMember(Name = ""test12"")"); diff --git a/ProtoAttributor.Tests/DataContracts/DataAttributeRewriterTests.cs b/ProtoAttributor.Tests/DataContracts/DataAttributeRewriterTests.cs index 3b22959..611ac15 100644 --- a/ProtoAttributor.Tests/DataContracts/DataAttributeRewriterTests.cs +++ b/ProtoAttributor.Tests/DataContracts/DataAttributeRewriterTests.cs @@ -91,5 +91,21 @@ public void RewritesAttributesWithCorrectOrderWhenAttributeExistsWithoutOrderPro output.Should().Contain("[DataMember(Order = 4)]"); _fixture.AssertOutputContainsCount(source, "[IgnoreDataMember]", 1); } + + [Fact] + public void RewritesEnumAttributesWhenFileHasDataMemeberIgnores() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithDataMemberIgnore.cs")); + var rewriter = new DataAttributeRewriter(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("System.Runtime.Serialization"); + output.Should().Contain("[DataContract]"); + _fixture.AssertOutputContainsCount(source, "[IgnoreDataMember]", 2); + _fixture.AssertOutputContainsCount(source, "[EnumMember]", 3); + } } } diff --git a/ProtoAttributor.Tests/Mocks/TestClassPlain.cs b/ProtoAttributor.Tests/Mocks/TestClassPlain.cs index 1cff18e..bca59e9 100644 --- a/ProtoAttributor.Tests/Mocks/TestClassPlain.cs +++ b/ProtoAttributor.Tests/Mocks/TestClassPlain.cs @@ -1,7 +1,10 @@ using System.ComponentModel.DataAnnotations; +using System.Runtime.Serialization; +using ProtoBuf; namespace ProtoAttributor.Tests.Mocks { + [DataContract] public class TestClassPlain { [Required] diff --git a/ProtoAttributor.Tests/Mocks/TestEnum.cs b/ProtoAttributor.Tests/Mocks/TestEnum.cs new file mode 100644 index 0000000..7714968 --- /dev/null +++ b/ProtoAttributor.Tests/Mocks/TestEnum.cs @@ -0,0 +1,13 @@ + +namespace ProtoAttributor.Tests.Mocks +{ + + public enum TestEnum + { + One, + Two, + Three, + Four, + Five + } +} diff --git a/ProtoAttributor.Tests/Mocks/TestEnumWithDataAttributes.cs b/ProtoAttributor.Tests/Mocks/TestEnumWithDataAttributes.cs new file mode 100644 index 0000000..ea6c217 --- /dev/null +++ b/ProtoAttributor.Tests/Mocks/TestEnumWithDataAttributes.cs @@ -0,0 +1,18 @@ + +using System.Runtime.Serialization; + +namespace ProtoAttributor.Tests.Mocks +{ + + [DataContract] + public enum TestEnumWithDataAttributes + { + [EnumMember] + One, + [EnumMember] + Two, + Three, + Four, + Five + } +} diff --git a/ProtoAttributor.Tests/Mocks/TestEnumWithDataMemberIgnore.cs b/ProtoAttributor.Tests/Mocks/TestEnumWithDataMemberIgnore.cs new file mode 100644 index 0000000..17513e9 --- /dev/null +++ b/ProtoAttributor.Tests/Mocks/TestEnumWithDataMemberIgnore.cs @@ -0,0 +1,17 @@ + +using System.Runtime.Serialization; + +namespace ProtoAttributor.Tests.Mocks +{ + [DataContract] + public enum TestEnumWithDataMemberIgnore + { + One, + Two, + [IgnoreDataMember] + Three, + Four, + [IgnoreDataMember] + Five + } +} diff --git a/ProtoAttributor.Tests/Mocks/TestEnumWithProtoAttributes.cs b/ProtoAttributor.Tests/Mocks/TestEnumWithProtoAttributes.cs new file mode 100644 index 0000000..2a64dcf --- /dev/null +++ b/ProtoAttributor.Tests/Mocks/TestEnumWithProtoAttributes.cs @@ -0,0 +1,17 @@ + +using ProtoBuf; + +namespace ProtoAttributor.Tests.Mocks +{ + [ProtoContract] + public enum TestEnumWithProtoAttributes + { + One, + [ProtoEnum] + Two, + [ProtoEnum] + Three, + Four, + Five + } +} diff --git a/ProtoAttributor.Tests/Mocks/TestEnumWithProtoIgnore.cs b/ProtoAttributor.Tests/Mocks/TestEnumWithProtoIgnore.cs new file mode 100644 index 0000000..64aaa6c --- /dev/null +++ b/ProtoAttributor.Tests/Mocks/TestEnumWithProtoIgnore.cs @@ -0,0 +1,17 @@ + +using ProtoBuf; + +namespace ProtoAttributor.Tests.Mocks +{ + [ProtoContract] + public enum TestEnumWithProtoIgnore + { + One, + [ProtoIgnore] + Two, + Three, + Four, + [ProtoIgnore] + Five + } +} diff --git a/ProtoAttributor.Tests/Mocks/TestRemoveAttributes.cs b/ProtoAttributor.Tests/Mocks/TestRemoveAttributes.cs index d3a6f46..ce67b24 100644 --- a/ProtoAttributor.Tests/Mocks/TestRemoveAttributes.cs +++ b/ProtoAttributor.Tests/Mocks/TestRemoveAttributes.cs @@ -25,4 +25,19 @@ public class TestRemoveAttributes public int MyProperty5 { get; set; } } + + [ProtoContract] + public enum TestRemoveEnumWithProtoAttributes + { + [ProtoEnum] + One, + [ProtoEnum] + Two, + [ProtoEnum] + Three, + [ProtoEnum] + Four, + [ProtoEnum] + Five + } } diff --git a/ProtoAttributor.Tests/Mocks/TestRemoveDataAttributes.cs b/ProtoAttributor.Tests/Mocks/TestRemoveDataAttributes.cs index 82e1bfc..2066d0b 100644 --- a/ProtoAttributor.Tests/Mocks/TestRemoveDataAttributes.cs +++ b/ProtoAttributor.Tests/Mocks/TestRemoveDataAttributes.cs @@ -28,4 +28,19 @@ public class TestRemoveDataAttributes [DataMember(Name = "test12")] public int MyProperty6 { get; set; } } + + [DataContract] + public enum TestRemoveEnumWithDataAttributes + { + [EnumMember] + One, + [EnumMember] + Two, + [EnumMember] + Three, + [EnumMember] + Four, + [EnumMember] + Five + } } diff --git a/ProtoAttributor.Tests/ProtoAttributor.Tests.csproj b/ProtoAttributor.Tests/ProtoAttributor.Tests.csproj index 02e937b..0b7eb34 100644 --- a/ProtoAttributor.Tests/ProtoAttributor.Tests.csproj +++ b/ProtoAttributor.Tests/ProtoAttributor.Tests.csproj @@ -30,6 +30,21 @@ Always + + Always + + + Always + + + Always + + + Always + + + Always + Always diff --git a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeAdderTests.cs b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeAdderTests.cs index 4a6b37d..6026307 100644 --- a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeAdderTests.cs +++ b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeAdderTests.cs @@ -154,5 +154,54 @@ public void AddsAttributesWithCorrectOrderWhenFileHasProtoIgnores() output.Should().Contain("[ProtoMember(16)]"); _fixture.AssertOutputContainsCount(source, "[ProtoIgnore]", 2); } + + [Fact] + public void AddsAttributesWhenFileIsEnum() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnum.cs")); + var rewriter = new ProtoAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("ProtoBuf"); + output.Should().Contain("[ProtoContract]"); + output.Should().Contain("[ProtoEnum]"); + _fixture.AssertOutputContainsCount(source, "[ProtoEnum]", 5); + } + + [Fact] + public void AddsAttributesWhenFileIsEnumWithExistingAttributes() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithProtoAttributes.cs")); + var rewriter = new ProtoAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("ProtoBuf"); + output.Should().Contain("[ProtoContract]"); + output.Should().Contain("[ProtoEnum]"); + _fixture.AssertOutputContainsCount(source, "[ProtoEnum]", 5); + } + + [Fact] + public void AddsAttributesWhenFileIsEnumWithIgnoreAttributes() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithProtoIgnore.cs")); + var rewriter = new ProtoAttributeAdder(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("ProtoBuf"); + output.Should().Contain("[ProtoContract]"); + output.Should().Contain("[ProtoEnum]"); + _fixture.AssertOutputContainsCount(source, "[ProtoEnum]", 3); + _fixture.AssertOutputContainsCount(source, "[ProtoIgnore]", 2); + } } } diff --git a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRemoverTests.cs b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRemoverTests.cs index 9a63796..9423281 100644 --- a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRemoverTests.cs +++ b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRemoverTests.cs @@ -15,7 +15,7 @@ public ProtoAttributeRemoverTests(TestFixure fixture) } [Fact] - public void AddsAttributesWithCorrectOrderWhenAttributesAlreadyExists() + public void RemovesProtoAttributesWhenAttributesAlreadyExists() { var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestRemoveAttributes.cs")); var rewriter = new ProtoAttributeRemover(); @@ -27,6 +27,7 @@ public void AddsAttributesWithCorrectOrderWhenAttributesAlreadyExists() output.Should().NotContain("ProtoBuf"); output.Should().NotContain("[ProtoContract]"); output.Should().NotContain("[ProtoInclude]"); + output.Should().NotContain("[ProtoEnum]"); output.Should().NotContain("[ProtoIgnore]"); output.Should().NotContain(@"[ProtoMember(1, Name=""Test"")]"); output.Should().NotContain("[ProtoMember(2)]"); diff --git a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRewriterTests.cs b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRewriterTests.cs index 27e394c..e997c4e 100644 --- a/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRewriterTests.cs +++ b/ProtoAttributor.Tests/ProtoContracts/ProtoAttributeRewriterTests.cs @@ -72,5 +72,21 @@ public void RewritesAttributesWithCorrectOrderWhenFileHasProtoIgnores() output.Should().Contain("[ProtoMember(4)]"); _fixture.AssertOutputContainsCount(source, "[ProtoIgnore]", 2); } + + [Fact] + public void RewritesEnumAttributesFileHasProtoIgnores() + { + var tree = CSharpSyntaxTree.ParseText(_fixture.LoadTestFile(@"./Mocks/TestEnumWithProtoIgnore.cs")); + var rewriter = new ProtoAttributeRewriter(); + var rewrittenRoot = rewriter.Visit(tree.GetRoot()); + + var output = rewrittenRoot.GetText().ToString(); + var source = output.Split(new string[] { " ", "\r\n" }, StringSplitOptions.RemoveEmptyEntries); + + output.Should().Contain("ProtoBuf"); + output.Should().Contain("[ProtoContract]"); + _fixture.AssertOutputContainsCount(source, "[ProtoIgnore]", 2); + _fixture.AssertOutputContainsCount(source, "[ProtoEnum]", 3); + } } } diff --git a/ProtoAttributor/Constants.cs b/ProtoAttributor/Constants.cs index 5fb82cd..d7dd68c 100644 --- a/ProtoAttributor/Constants.cs +++ b/ProtoAttributor/Constants.cs @@ -7,6 +7,8 @@ public static class Proto public const string PROPERTY_ATTRIBUTE_NAME = "ProtoMember"; public const string PROPERTY_IGNORE_ATTRIBUTE_NAME = "ProtoIgnore"; public const string CLASS_ATTRIBUTE_NAME = "ProtoContract"; + public const string ENUM_ATTRIBUTE_NAME = "ProtoContract"; + public const string ENUM_MEMBER_NAME = "ProtoEnum"; public const string USING_STATEMENT = "ProtoBuf"; public const string BASE_PROP_NAME = "Proto"; } @@ -16,6 +18,8 @@ public static class Data public const string PROPERTY_ATTRIBUTE_NAME = "DataMember"; public const string PROPERTY_IGNORE_ATTRIBUTE_NAME = "IgnoreDataMember"; public const string CLASS_ATTRIBUTE_NAME = "DataContract"; + public const string ENUM_ATTRIBUTE_NAME = "DataContract"; + public const string ENUM_MEMBER_NAME = "EnumMember"; public const string USING_STATEMENT = "System.Runtime.Serialization"; public const string BASE_PROP_NAME = "Data"; public const string BASE_PROPIGNORE_NAME = "IgnoreData"; diff --git a/ProtoAttributor/Parsers/DataContracts/BaseDataRewriter.cs b/ProtoAttributor/Parsers/DataContracts/BaseDataRewriter.cs index 46f37ff..2f86a34 100644 --- a/ProtoAttributor/Parsers/DataContracts/BaseDataRewriter.cs +++ b/ProtoAttributor/Parsers/DataContracts/BaseDataRewriter.cs @@ -41,6 +41,28 @@ public override SyntaxNode VisitCompilationUnit(CompilationUnitSyntax node) return base.VisitCompilationUnit(node); } + public override SyntaxNode VisitEnumDeclaration(EnumDeclarationSyntax node) + { + //each class needs to restat with the + _startIndex = CalculateStartingIndex(node); + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Data.ENUM_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Data.ENUM_ATTRIBUTE_NAME); + var attribute = SyntaxFactory.Attribute(name); + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + } + + return base.VisitEnumDeclaration(node); + } + public override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax node) { //each class needs to restat with the diff --git a/ProtoAttributor/Parsers/DataContracts/DataAttributeAdder.cs b/ProtoAttributor/Parsers/DataContracts/DataAttributeAdder.cs index 0b6c157..ddb27d1 100644 --- a/ProtoAttributor/Parsers/DataContracts/DataAttributeAdder.cs +++ b/ProtoAttributor/Parsers/DataContracts/DataAttributeAdder.cs @@ -17,6 +17,26 @@ public override int CalculateStartingIndex(SyntaxNode node) { return _dataReader.GetDataMemberNextId(node); } + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Data.ENUM_MEMBER_NAME, Constants.Data.PROPERTY_IGNORE_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Data.ENUM_MEMBER_NAME); + var attribute = SyntaxFactory.Attribute(name); //EnumMember + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + _startIndex++; + + } + return base.VisitEnumMemberDeclaration(node); + } public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax node) { diff --git a/ProtoAttributor/Parsers/DataContracts/DataAttributeRemover.cs b/ProtoAttributor/Parsers/DataContracts/DataAttributeRemover.cs index 8023899..350ab40 100644 --- a/ProtoAttributor/Parsers/DataContracts/DataAttributeRemover.cs +++ b/ProtoAttributor/Parsers/DataContracts/DataAttributeRemover.cs @@ -26,6 +26,30 @@ public override SyntaxNode VisitCompilationUnit(CompilationUnitSyntax node) return base.VisitCompilationUnit(node); } + public override SyntaxNode VisitEnumDeclaration(EnumDeclarationSyntax node) + { + if (node.AttributeLists.Count > 0) + { + var newAttributeLists = new SyntaxList(); + foreach (var attributeList in node.AttributeLists) + { + var nodesToRemove = attributeList.Attributes.Where(attribute => NodeHelper.AttributeNameContains(attribute, Constants.Data.BASE_PROP_NAME, Constants.Data.BASE_KNOWN_TYPE_NAME)).ToArray(); + + // If the lists are the same length, we are removing all attributes and can just avoid populating newAttributes. + if (nodesToRemove.Length != attributeList.Attributes.Count) + { + var newAttribute = (AttributeListSyntax)VisitAttributeList(attributeList.RemoveNodes(nodesToRemove, SyntaxRemoveOptions.KeepNoTrivia)); + newAttributeLists = newAttributeLists.Add(newAttribute); + } + } + var leadTriv = node.GetLeadingTrivia(); + node = node.WithAttributeLists(newAttributeLists); + node = node.WithLeadingTrivia(leadTriv); + } + + return base.VisitEnumDeclaration(node); + } + public override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax node) { if (node.AttributeLists.Count > 0) @@ -73,5 +97,29 @@ public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax no return base.VisitPropertyDeclaration(node); } + + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + if (node.AttributeLists.Count > 0) + { + var newAttributeLists = new SyntaxList(); + foreach (var attributeList in node.AttributeLists) + { + var nodesToRemove = attributeList.Attributes.Where(attribute => NodeHelper.AttributeNameContains(attribute, Constants.Data.BASE_PROP_NAME, Constants.Data.BASE_PROPIGNORE_NAME, Constants.Data.ENUM_MEMBER_NAME)).ToArray(); + + // If the lists are the same length, we are removing all attributes and can just avoid populating newAttributes. + if (nodesToRemove.Length != attributeList.Attributes.Count) + { + var newAttribute = (AttributeListSyntax)VisitAttributeList(attributeList.RemoveNodes(nodesToRemove, SyntaxRemoveOptions.KeepNoTrivia)); + newAttributeLists = newAttributeLists.Add(newAttribute); + } + } + var leadTriv = node.GetLeadingTrivia(); + node = node.WithAttributeLists(newAttributeLists); + node = node.WithLeadingTrivia(leadTriv); + } + + return base.VisitEnumMemberDeclaration(node); + } } } diff --git a/ProtoAttributor/Parsers/DataContracts/DataAttributeRewriter.cs b/ProtoAttributor/Parsers/DataContracts/DataAttributeRewriter.cs index ac07276..e5fe536 100644 --- a/ProtoAttributor/Parsers/DataContracts/DataAttributeRewriter.cs +++ b/ProtoAttributor/Parsers/DataContracts/DataAttributeRewriter.cs @@ -14,6 +14,27 @@ public override int CalculateStartingIndex(SyntaxNode node) return 1; } + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Data.ENUM_MEMBER_NAME, Constants.Data.PROPERTY_IGNORE_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Data.ENUM_MEMBER_NAME); + var attribute = SyntaxFactory.Attribute(name); //EnumMember + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + + _startIndex++; + } + return base.VisitEnumMemberDeclaration(node); + } + public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax node) { var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Data.PROPERTY_ATTRIBUTE_NAME, Constants.Data.PROPERTY_IGNORE_ATTRIBUTE_NAME); diff --git a/ProtoAttributor/Parsers/ProtoContracts/BaseProtoRewriter.cs b/ProtoAttributor/Parsers/ProtoContracts/BaseProtoRewriter.cs index 2d6c2a2..c13bc15 100644 --- a/ProtoAttributor/Parsers/ProtoContracts/BaseProtoRewriter.cs +++ b/ProtoAttributor/Parsers/ProtoContracts/BaseProtoRewriter.cs @@ -42,6 +42,28 @@ public override SyntaxNode VisitCompilationUnit(CompilationUnitSyntax node) return base.VisitCompilationUnit(node); } + public override SyntaxNode VisitEnumDeclaration(EnumDeclarationSyntax node) + { + //each class needs to restat with the + _startIndex = CalculateStartingIndex(node); + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Proto.ENUM_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Proto.ENUM_ATTRIBUTE_NAME); + var attribute = SyntaxFactory.Attribute(name); + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + } + + return base.VisitEnumDeclaration(node); + } + public override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax node) { //each class needs to restat with the diff --git a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeAdder.cs b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeAdder.cs index 87aa97f..345a388 100644 --- a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeAdder.cs +++ b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeAdder.cs @@ -20,6 +20,27 @@ public override int CalculateStartingIndex(SyntaxNode node) return _protoReader.GetProtoNextId(node); } + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Proto.ENUM_MEMBER_NAME, Constants.Proto.PROPERTY_IGNORE_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Proto.ENUM_MEMBER_NAME); + var attribute = SyntaxFactory.Attribute(name); //ProtoEnum() + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + _startIndex++; + } + + return base.VisitEnumMemberDeclaration(node); + } + public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax node) { var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Proto.PROPERTY_ATTRIBUTE_NAME, Constants.Proto.PROPERTY_IGNORE_ATTRIBUTE_NAME); diff --git a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRemover.cs b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRemover.cs index 04985ca..896be93 100644 --- a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRemover.cs +++ b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRemover.cs @@ -27,6 +27,29 @@ public override SyntaxNode VisitCompilationUnit(CompilationUnitSyntax node) return base.VisitCompilationUnit(node); } + public override SyntaxNode VisitEnumDeclaration(EnumDeclarationSyntax node) + { + if (node.AttributeLists.Count > 0) + { + var newAttributeLists = new SyntaxList(); + foreach (var attributeList in node.AttributeLists) + { + var nodesToRemove = attributeList.Attributes.Where(attribute => NodeHelper.AttributeNameMatches(attribute, Constants.Proto.BASE_PROP_NAME)).ToArray(); + + // If the lists are the same length, we are removing all attributes and can just avoid populating newAttributes. + if (nodesToRemove.Length != attributeList.Attributes.Count) + { + var newAttribute = (AttributeListSyntax)VisitAttributeList(attributeList.RemoveNodes(nodesToRemove, SyntaxRemoveOptions.KeepNoTrivia)); + newAttributeLists = newAttributeLists.Add(newAttribute); + } + } + var leadTriv = node.GetLeadingTrivia(); + node = node.WithAttributeLists(newAttributeLists); + node = node.WithLeadingTrivia(leadTriv); + } + return base.VisitEnumDeclaration(node); + } + public override SyntaxNode VisitClassDeclaration(ClassDeclarationSyntax node) { if (node.AttributeLists.Count > 0) @@ -74,5 +97,29 @@ public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax no return base.VisitPropertyDeclaration(node); } + + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + if (node.AttributeLists.Count > 0) + { + var newAttributeLists = new SyntaxList(); + foreach (var attributeList in node.AttributeLists) + { + var nodesToRemove = attributeList.Attributes.Where(attribute => NodeHelper.AttributeNameMatches(attribute, Constants.Proto.BASE_PROP_NAME)).ToArray(); + + // If the lists are the same length, we are removing all attributes and can just avoid populating newAttributes. + if (nodesToRemove.Length != attributeList.Attributes.Count) + { + var newAttribute = (AttributeListSyntax)VisitAttributeList(attributeList.RemoveNodes(nodesToRemove, SyntaxRemoveOptions.KeepNoTrivia)); + newAttributeLists = newAttributeLists.Add(newAttribute); + } + } + var leadTriv = node.GetLeadingTrivia(); + node = node.WithAttributeLists(newAttributeLists); + node = node.WithLeadingTrivia(leadTriv); + + } + return base.VisitEnumMemberDeclaration(node); + } } } diff --git a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRewriter.cs b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRewriter.cs index dad5a02..461e549 100644 --- a/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRewriter.cs +++ b/ProtoAttributor/Parsers/ProtoContracts/ProtoAttributeRewriter.cs @@ -15,6 +15,25 @@ public override int CalculateStartingIndex(SyntaxNode node) return 1; } + public override SyntaxNode VisitEnumMemberDeclaration(EnumMemberDeclarationSyntax node) + { + var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Proto.ENUM_MEMBER_NAME, Constants.Proto.PROPERTY_IGNORE_ATTRIBUTE_NAME); + + if (!hasMatch) + { + var name = SyntaxFactory.ParseName(Constants.Proto.ENUM_MEMBER_NAME); + var attribute = SyntaxFactory.Attribute(name); //ProtoEnum() + + node = TriviaMaintainer.Apply(node, (innerNode, wp) => + { + var newAttributes = BuildAttribute(attribute, innerNode.AttributeLists, wp); + + return innerNode.WithAttributeLists(newAttributes).WithAdditionalAnnotations(Formatter.Annotation); + }); + } + return base.VisitEnumMemberDeclaration(node); + } + public override SyntaxNode VisitPropertyDeclaration(PropertyDeclarationSyntax node) { var hasMatch = NodeHelper.HasMatch(node.AttributeLists, Constants.Proto.PROPERTY_ATTRIBUTE_NAME, Constants.Proto.PROPERTY_IGNORE_ATTRIBUTE_NAME); diff --git a/ProtoAttributor/Parsers/TriviaMaintainer.cs b/ProtoAttributor/Parsers/TriviaMaintainer.cs index 221bbdb..15a205c 100644 --- a/ProtoAttributor/Parsers/TriviaMaintainer.cs +++ b/ProtoAttributor/Parsers/TriviaMaintainer.cs @@ -45,5 +45,43 @@ public static ClassDeclarationSyntax Apply(ClassDeclarationSyntax node, Func builder) + { + var leadingTrivia = node.GetLeadingTrivia(); + var trailingTrivia = node.GetTrailingTrivia(); + + node = node.WithoutLeadingTrivia(); + node = node.WithoutTrailingTrivia(); + + var wp = leadingTrivia.FirstOrDefault(w => w.Kind() == SyntaxKind.WhitespaceTrivia); + + node = builder?.Invoke(node, wp); + + node = node.WithLeadingTrivia(leadingTrivia); + + node = node.WithTrailingTrivia(trailingTrivia); + + return node; + } + + public static EnumMemberDeclarationSyntax Apply(EnumMemberDeclarationSyntax node, Func builder) + { + var leadingTrivia = node.GetLeadingTrivia(); + var trailingTrivia = node.GetTrailingTrivia(); + + node = node.WithoutLeadingTrivia(); + node = node.WithoutTrailingTrivia(); + + var wp = leadingTrivia.FirstOrDefault(w => w.Kind() == SyntaxKind.WhitespaceTrivia); + + node = builder?.Invoke(node, wp); + + node = node.WithLeadingTrivia(leadingTrivia); + + node = node.WithTrailingTrivia(trailingTrivia); + + return node; + } } } diff --git a/ProtoAttributor/source.extension.vsixmanifest b/ProtoAttributor/source.extension.vsixmanifest index 9852ea5..3d1e262 100644 --- a/ProtoAttributor/source.extension.vsixmanifest +++ b/ProtoAttributor/source.extension.vsixmanifest @@ -1,7 +1,7 @@ - + ProtoAttributor ProtoAttributor is an open source Visual Studio extension that can manage the appropriate attributes on a class to support ProtoBuf. It currently supports ProtoContract, ProtoMember, ProtoIgnore, DataContract, DataMember, IgnoreDataMemeber attributes. https://github.com/d1820/proto-attributor