Allow transforms to remove enum members

This fixes a bug that was preventing enum members from being removable from transforms, such as removeShapesByTag.
sugmanue · Oct 11, 2022 · b9d5a49 · b9d5a49
1 parent 23a11e5
commit b9d5a49
Show file tree

Hide file tree

Showing 4 changed files with 69 additions and 2 deletions.
diff --git a/...y-build/src/test/java/software/amazon/smithy/build/transforms/ExcludeShapesByTagTest.java b/...y-build/src/test/java/software/amazon/smithy/build/transforms/ExcludeShapesByTagTest.java
@@ -19,11 +19,15 @@
 import static org.hamcrest.Matchers.is;
 import static org.hamcrest.Matchers.not;
 
+import java.nio.file.Paths;
 import java.util.Optional;
 import org.junit.jupiter.api.Test;
 import software.amazon.smithy.build.TransformContext;
 import software.amazon.smithy.model.Model;
 import software.amazon.smithy.model.node.Node;
+import software.amazon.smithy.model.shapes.EnumShape;
+import software.amazon.smithy.model.shapes.IntEnumShape;
+import software.amazon.smithy.model.shapes.ShapeId;
 import software.amazon.smithy.model.shapes.StringShape;
 import software.amazon.smithy.model.traits.TagsTrait;
 
@@ -51,4 +55,23 @@ public void removesTraitsNotInList() {
         assertThat(result.getShape(stringA.getId()), is(Optional.empty()));
         assertThat(result.getShape(stringB.getId()), not(Optional.empty()));
     }
+
+    @Test
+    public void filtersMembers() throws Exception {
+        Model model = Model.assembler()
+                .addImport(Paths.get(getClass().getResource("filter-by-tags.smithy").toURI()))
+                .assemble()
+                .unwrap();
+        TransformContext context = TransformContext.builder()
+                .model(model)
+                .settings(Node.objectNode().withMember("tags", Node.fromStrings("filter")))
+                .build();
+        Model result = new ExcludeShapesByTag().transform(context);
+
+        EnumShape foo = result.expectShape(ShapeId.from("smithy.example#Foo"), EnumShape.class);
+        assertThat(foo.members().size(), is(1));
+
+        IntEnumShape bar = result.expectShape(ShapeId.from("smithy.example#Bar"), IntEnumShape.class);
+        assertThat(bar.members().size(), is(1));
+    }
 }
diff --git a/...hy-build/src/test/resources/software/amazon/smithy/build/transforms/filter-by-tags.smithy b/...hy-build/src/test/resources/software/amazon/smithy/build/transforms/filter-by-tags.smithy
@@ -0,0 +1,21 @@
+$version: "2.0"
+
+namespace smithy.example
+
+enum Foo {
+    @tags(["filter"])
+    FILTER
+
+    @tags(["keep"])
+    KEEP
+}
+
+intEnum Bar {
+    @tags(["filter"])
+    @enumValue(1)
+    FILTER
+
+    @tags(["keep"])
+    @enumValue(2)
+    KEEP
+}
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/transform/FilterShapes.java b/smithy-model/src/main/java/software/amazon/smithy/model/transform/FilterShapes.java
@@ -51,7 +51,10 @@ Model transform(ModelTransformer transformer, Model model) {
 
     private static boolean canFilterShape(Model model, Shape shape) {
         return !shape.isMemberShape() || model.getShape(shape.asMemberShape().get().getContainer())
-                .filter(container -> container.isStructureShape() || container.isUnionShape())
+                .filter(container -> container.isStructureShape()
+                        || container.isUnionShape()
+                        || container.isEnumShape()
+                        || container.isIntEnumShape())
                 .isPresent();
     }
 }
diff --git a/...in/java/software/amazon/smithy/model/transform/plugins/CleanStructureAndUnionMembers.java b/...in/java/software/amazon/smithy/model/transform/plugins/CleanStructureAndUnionMembers.java
@@ -28,6 +28,8 @@
 import java.util.function.Function;
 import java.util.stream.Collectors;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.EnumShape;
+import software.amazon.smithy.model.shapes.IntEnumShape;
 import software.amazon.smithy.model.shapes.MemberShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeId;
@@ -39,7 +41,7 @@
 import software.amazon.smithy.utils.Pair;
 
 /**
- * Cleans up structure and union shapes after shapes are removed.
+ * Cleans up structure, union, enum, and intEnum shapes after shapes are removed.
  *
  * <ul>
  *     <li>Ensures that structure and union shapes are updated to no
@@ -60,9 +62,27 @@ public Model onRemove(ModelTransformer transformer, Collection<Shape> removed, M
     private Model removeMembersFromContainers(ModelTransformer transformer, Collection<Shape> removed, Model model) {
         List<Shape> replacements = new ArrayList<>(getStructureReplacements(model, removed));
         replacements.addAll(getUnionReplacements(model, removed));
+        replacements.addAll(getEnumReplacements(model, removed));
+        replacements.addAll(getIntEnumReplacements(model, removed));
         return transformer.replaceShapes(model, replacements);
     }
 
+    private Collection<Shape> getEnumReplacements(Model model, Collection<Shape> removed) {
+        return createUpdatedShapes(model, removed, Shape::asEnumShape, entry -> {
+            EnumShape.Builder builder = entry.getKey().toBuilder();
+            entry.getValue().forEach(member -> builder.removeMember(member.getMemberName()));
+            return builder.build();
+        });
+    }
+
+    private Collection<Shape> getIntEnumReplacements(Model model, Collection<Shape> removed) {
+        return createUpdatedShapes(model, removed, Shape::asIntEnumShape, entry -> {
+            IntEnumShape.Builder builder = entry.getKey().toBuilder();
+            entry.getValue().forEach(member -> builder.removeMember(member.getMemberName()));
+            return builder.build();
+        });
+    }
+
     private Collection<Shape> getStructureReplacements(Model model, Collection<Shape> removed) {
         return createUpdatedShapes(model, removed, Shape::asStructureShape, entry -> {
             StructureShape.Builder builder = entry.getKey().toBuilder();