Add filtering predicates to idl serializer

smithy-lang · JordonPhillips · Feb 27, 2020 · Feb 21, 2020 · Feb 21, 2020 · Feb 24, 2020
commit e8097964862335d775b6628fef9c4b7fbdbec746
diff --git a/smithy-model/src/main/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializer.java b/smithy-model/src/main/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializer.java
@@ -27,6 +27,7 @@
 import java.util.Objects;
 import java.util.Set;
 import java.util.function.Function;
+import java.util.function.Predicate;
 import java.util.regex.Pattern;
 import java.util.stream.Collectors;
 import software.amazon.smithy.model.Model;
@@ -52,9 +53,15 @@
  * Serializes a {@link Model} into a set of Smithy IDL files.
  */
 public final class SmithyIdlModelSerializer {
+    private final Predicate<String> metadataFilter;
+    private final Predicate<Shape> shapeFilter;
+    private final Predicate<Trait> traitFilter;
     private final Function<Shape, Path> shapePlacer;
 
     private SmithyIdlModelSerializer(Builder builder) {
+        metadataFilter = builder.metadataFilter;
+        shapeFilter = builder.shapeFilter.and(FunctionalUtils.not(Prelude::isPreludeShape));
+        traitFilter = builder.traitFilter;
         shapePlacer = builder.shapePlacer;
     }
 
@@ -67,7 +74,8 @@ private SmithyIdlModelSerializer(Builder builder) {
      */
     public Map<Path, String> serialize(Model model) {
         return model.shapes()
-                .filter(FunctionalUtils.not(Prelude::isPreludeShape))
+                .filter(FunctionalUtils.not(Shape::isMemberShape))
+                .filter(shapeFilter)
                 .collect(Collectors.groupingBy(shapePlacer)).entrySet().stream()
                 .collect(Collectors.toMap(Map.Entry::getKey, entry -> serialize(model, entry.getValue())));
     }
@@ -85,7 +93,7 @@ private String serialize(Model fullModel, Collection<Shape> shapes) {
         SmithyCodeWriter codeWriter = new SmithyCodeWriter(namespace, fullModel);
         NodeSerializer nodeSerializer = new NodeSerializer(codeWriter, fullModel);
 
-        ShapeSerializer shapeSerializer = new ShapeSerializer(codeWriter, nodeSerializer, fullModel);
+        ShapeSerializer shapeSerializer = new ShapeSerializer(codeWriter, nodeSerializer, traitFilter, fullModel);
         shapes.stream()
                 .filter(FunctionalUtils.not(Shape::isMemberShape))
                 .sorted(new ShapeComparator())
@@ -104,6 +112,7 @@ private String serializeHeader(String namespace, Model fullModel) {
         // but if they're separated out then each file will still have all the context.
         fullModel.getMetadata().entrySet().stream()
                 .sorted(Map.Entry.comparingByKey(String.CASE_INSENSITIVE_ORDER))
+                .filter(entry -> metadataFilter.test(entry.getKey()))
                 .forEach(entry -> {
                     codeWriter.trimTrailingSpaces(false)
                             .writeInline("metadata $M = ", entry.getKey())
@@ -190,10 +199,48 @@ private int compareCaseInsensitive(Shape s1, Shape s2) {
      * Builder used to create {@link SmithyIdlModelSerializer}.
      */
     public static final class Builder implements SmithyBuilder<SmithyIdlModelSerializer> {
+        private Predicate<String> metadataFilter = pair -> true;
+        private Predicate<Shape> shapeFilter = shape -> true;
+        private Predicate<Trait> traitFilter = trait -> true;
         private Function<Shape, Path> shapePlacer = SmithyIdlModelSerializer::placeShapesByNamespace;
 
         public Builder() {}
 
+        /**
+         * Predicate that determines if a metadata is serialized.
+         * @param metadataFilter Predicate that accepts a metadata key.
+         * @return Returns the builder.
+         */
+        public Builder metadataFilter(Predicate<String> metadataFilter) {
+            this.metadataFilter = Objects.requireNonNull(metadataFilter);
+            return this;
+        }
+
+        /**
+         * Predicate that determines if a shape and its traits are serialized.
+         * @param shapeFilter Predicate that accepts a shape.
+         * @return Returns the builder.
+         */
+        public Builder shapeFilter(Predicate<Shape> shapeFilter) {
+            this.shapeFilter = Objects.requireNonNull(shapeFilter);
+            return this;
+        }
+
+        /**
+         * Sets a predicate that can be used to filter trait values from
+         * appearing in the serialized model.
+         *
+         * <p>Note that this does not filter out trait definitions. It only filters
+         * out instances of traits from being serialized on shapes.
+         *
+         * @param traitFilter Predicate that filters out trait definitions.
+         * @return Returns the builder.
+         */
+        public Builder traitFilter(Predicate<Trait> traitFilter) {
+            this.traitFilter = traitFilter;
+            return this;
+        }
+
         /**
          * Function that determines what output file a shape should go in.
          *
@@ -218,11 +265,18 @@ public SmithyIdlModelSerializer build() {
     private static final class ShapeSerializer extends ShapeVisitor.Default<Void> {
         private final SmithyCodeWriter codeWriter;
         private final NodeSerializer nodeSerializer;
+        private final Predicate<Trait> traitFilter;
         private final Model model;
 
-        ShapeSerializer(SmithyCodeWriter codeWriter, NodeSerializer nodeSerializer, Model model) {
+        ShapeSerializer(
+                SmithyCodeWriter codeWriter,
+                NodeSerializer nodeSerializer,
+                Predicate<Trait> traitFilter,
+                Model model
+        ) {
             this.codeWriter = codeWriter;
             this.nodeSerializer = nodeSerializer;
+            this.traitFilter = traitFilter;
             this.model = model;
         }
 
@@ -251,9 +305,10 @@ private void shapeWithMembers(Shape shape, List<MemberShape> members) {
 
         private void serializeTraits(Shape shape) {
             // The documentation trait always needs to be serialized first since it uses special syntax.
-            shape.getTrait(DocumentationTrait.class).ifPresent(this::serializeDocumentationTrait);
+            shape.getTrait(DocumentationTrait.class).filter(traitFilter).ifPresent(this::serializeDocumentationTrait);
             shape.getAllTraits().values().stream()
                     .filter(trait -> !(trait instanceof DocumentationTrait))
+                    .filter(traitFilter)
                     .sorted(Comparator.comparing(trait -> trait.toShapeId().toString(), String.CASE_INSENSITIVE_ORDER))
                     .forEach(this::serializeTrait);
         }

diff --git a/...model/src/test/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializerTest.java b/...model/src/test/java/software/amazon/smithy/model/shapes/SmithyIdlModelSerializerTest.java
@@ -1,7 +1,11 @@
 package software.amazon.smithy.model.shapes;
 
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.equalTo;
+import static org.hamcrest.Matchers.aMapWithSize;
+import static org.hamcrest.Matchers.hasKey;
+import static org.hamcrest.Matchers.not;
 
 import java.io.IOException;
 import java.nio.file.Files;
@@ -13,6 +17,8 @@
 import org.junit.jupiter.api.Test;
 import org.junit.jupiter.api.TestFactory;
 import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.traits.DocumentationTrait;
+import software.amazon.smithy.model.traits.RequiredTrait;
 import software.amazon.smithy.utils.IoUtils;
 
 public class SmithyIdlModelSerializerTest {
@@ -51,4 +57,66 @@ public void multipleNamespacesGenerateMultipleFiles() {
             assertThat(generated, equalTo(IoUtils.readUtf8File(expectedPath)));
         });
     }
+
+    @Test
+    public void filtersShapes() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idl-serialization/test-model.json"))
+                .assemble()
+                .unwrap();
+        SmithyIdlModelSerializer serializer = SmithyIdlModelSerializer.builder()
+                .shapeFilter(shape -> shape.getId().getNamespace().equals("ns.structures"))
+                .build();
+        Map<Path, String> serialized = serializer.serialize(model);
+
+        assertThat(serialized, aMapWithSize(1));
+        assertThat(serialized, hasKey(Paths.get("ns.structures.smithy")));
+        assertThat(serialized.get(Paths.get("ns.structures.smithy")),
+                containsString("namespace ns.structures"));
+    }
+
+    @Test
+    public void filtersMetadata() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idl-serialization/test-model.json"))
+                .assemble()
+                .unwrap();
+        SmithyIdlModelSerializer serializer = SmithyIdlModelSerializer.builder()
+                .metadataFilter(key -> false)
+                .build();
+        Map<Path, String> serialized = serializer.serialize(model);
+        for (String output : serialized.values()) {
+            assertThat(output, not(containsString("metadata")));
+        }
+    }
+
+    @Test
+    public void filtersTraits() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idl-serialization/test-model.json"))
+                .assemble()
+                .unwrap();
+        SmithyIdlModelSerializer serializer = SmithyIdlModelSerializer.builder()
+                .traitFilter(trait -> !(trait instanceof RequiredTrait))
+                .build();
+        Map<Path, String> serialized = serializer.serialize(model);
+        for (String output : serialized.values()) {
+            assertThat(output, not(containsString("@required")));
+        }
+    }
+
+    @Test
+    public void filtersDocumentationTrait() {
+        Model model = Model.assembler()
+                .addImport(getClass().getResource("idl-serialization/test-model.json"))
+                .assemble()
+                .unwrap();
+        SmithyIdlModelSerializer serializer = SmithyIdlModelSerializer.builder()
+                .traitFilter(trait -> !(trait instanceof DocumentationTrait))
+                .build();
+        Map<Path, String> serialized = serializer.serialize(model);
+        for (String output : serialized.values()) {
+            assertThat(output, not(containsString("/// ")));
+        }
+    }
 }
diff --git a/.../src/test/resources/software/amazon/smithy/model/shapes/idl-serialization/test-model.json b/.../src/test/resources/software/amazon/smithy/model/shapes/idl-serialization/test-model.json
@@ -0,0 +1,44 @@
+{
+    "smithy": "0.5.0",
+    "metadata": {
+        "shared": true
+    },
+    "shapes": {
+        "ns.structures#Structure": {
+            "type": "structure",
+            "members": {
+                "listMember": {
+                    "target": "ns.primitives#StringList",
+                    "traits": {
+                        "smithy.api#documentation": "A list member"
+                    }
+                },
+                "stringMember": {
+                    "target": "ns.primitives#String",
+                    "traits": {
+                        "smithy.api#required": true,
+                        "smithy.api#documentation": "A string member"
+                    }
+                }
+            },
+            "traits": {
+                "smithy.api#documentation": "A structure shape"
+            }
+        },
+        "ns.primitives#String": {
+            "type": "string",
+            "traits": {
+                "smithy.api#documentation": "A string shape"
+            }
+        },
+        "ns.primitives#StringList": {
+            "type": "list",
+            "member": {
+                "target": "ns.primitives#String"
+            },
+            "traits": {
+                "smithy.api#documentation": "A list shape"
+            }
+        }
+    }
+}