Merge pull request #3 from d1820/add-enum-support

added enum support
d1820 · Feb 26, 2021 · 28b108b · 28b108b
2 parents 1cd53b6 + e2d6363
commit 28b108b
Show file tree

Hide file tree

Showing 26 changed files with 527 additions and 2 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -30,6 +30,21 @@
     <Compile Update="Mocks\TestCodeWithDataMemberAttributes.cs">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </Compile>
+    <Compile Update="Mocks\TestEnum.cs">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Compile>
+    <Compile Update="Mocks\TestEnumWithDataAttributes.cs">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Compile>
+    <Compile Update="Mocks\TestEnumWithDataMemberIgnore.cs">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Compile>
+    <Compile Update="Mocks\TestEnumWithProtoAttributes.cs">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Compile>
+    <Compile Update="Mocks\TestEnumWithProtoIgnore.cs">
+      <CopyToOutputDirectory>Always</CopyToOutputDirectory>
+    </Compile>
     <Compile Update="Mocks\TestMissingOrderPropertyAndWeirdSpacing.cs">
       <CopyToOutputDirectory>Always</CopyToOutputDirectory>
     </Compile>

diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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