Enum member removal (again)

smithy-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));
+    }
 }

smithy-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
+}

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();
     }
 }

smithy-model/src/main/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();

smithy-model/src/test/java/software/amazon/smithy/model/transform/RemoveShapesTest.java

@@ -21,6 +21,7 @@
 import static org.hamcrest.Matchers.hasKey;
 import static org.hamcrest.Matchers.not;
 import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertThrows;
 
 import java.util.ArrayList;
 import java.util.Arrays;
@@ -34,7 +35,10 @@
 import org.junit.jupiter.params.ParameterizedTest;
 import org.junit.jupiter.params.provider.MethodSource;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.SourceException;
 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.ListShape;
 import software.amazon.smithy.model.shapes.MapShape;
 import software.amazon.smithy.model.shapes.MemberShape;
@@ -103,6 +107,56 @@ public void removesMapMembersWhenRemoved() {
         assertContainerMembersAreRemoved(container, Arrays.asList(key, value));
     }
 
+    @Test
+    public void removesEnumMembersWhenRemoved() {
+        EnumShape container = EnumShape.builder()
+                .id(ShapeId.from("ns.foo#Enum"))
+                .addMember("foo", "foo")
+                .build();
+        assertContainerMembersAreRemoved(container, new ArrayList<>(container.members()));
+    }
+
+    @Test
+    public void notAllEnumMembersCanBeRemoved() {
+        EnumShape container = EnumShape.builder()
+                .id(ShapeId.from("ns.foo#Enum"))
+                .addMember("foo", "foo")
+                .build();
+
+        Model.Builder builder = Model.builder()
+                .addShape(container)
+                .addShape(StringShape.builder().id(STRING_TARGET).build());
+        container.members().forEach(builder::addShape);
+        Model model = builder.build();
+        ModelTransformer transformer = ModelTransformer.create();
+        assertThrows(SourceException.class, () -> transformer.removeShapes(model, container.members()));
+    }
+
+    @Test
+    public void removesIntEnumMembersWhenRemoved() {
+        IntEnumShape container = IntEnumShape.builder()
+                .id(ShapeId.from("ns.foo#Enum"))
+                .addMember("foo", 1)
+                .build();
+        assertContainerMembersAreRemoved(container, new ArrayList<>(container.members()));
+    }
+
+    @Test
+    public void notAllIntEnumMembersCanBeRemoved() {
+        IntEnumShape container = IntEnumShape.builder()
+                .id(ShapeId.from("ns.foo#Enum"))
+                .addMember("foo", 1)
+                .build();
+
+        Model.Builder builder = Model.builder()
+                .addShape(container)
+                .addShape(StringShape.builder().id(STRING_TARGET).build());
+        container.members().forEach(builder::addShape);
+        Model model = builder.build();
+        ModelTransformer transformer = ModelTransformer.create();
+        assertThrows(SourceException.class, () -> transformer.removeShapes(model, container.members()));
+    }
+
     @Test
     public void removesStructureMembersWhenRemoved() {
         MemberShape a = MemberShape.builder()