Add transformer that removes trait definitions

docs/source/1.0/guides/building-models/build-config.rst

@@ -862,6 +862,54 @@ key is not in the provided ``keys`` list.
             }
         }
 
+.. _removeTraitDefinitions-transform:
+
+removeTraitDefinitions
+----------------------
+
+Removes trait definitions from the model, but leaves the instances of traits
+intact on any shapes.
+
+You can *export* trait definitions by applying specific tags to the trait
+definition and adding the list of export tags in the ``exportTagged`` argument.
+
+.. list-table::
+    :header-rows: 1
+    :widths: 10 20 70
+
+    * - Property
+      - Type
+      - Description
+    * - exportTagged
+      - ``[string]``
+      - The set of tags that, if found on a trait definition, forces the trait
+        to be retained in the transformed model.
+
+The following example removes trait definitions but keeps the instances of the
+trait intact on shapes in the model:
+
+.. tabs::
+
+    .. code-tab:: json
+
+        {
+            "version": "1.0",
+            "projections": {
+                "exampleProjection": {
+                    "transforms": [
+                        {
+                            "name": "removeTraitDefinitions",
+                            "args": {
+                                "exportTagged": [
+                                    "export-tag1",
+                                    "another-export-tag"
+                                ]
+                            }
+                        }
+                    ]
+                }
+            }
+        }
 
 .. _removeUnusedShapes-transform:
 

smithy-build/src/main/java/software/amazon/smithy/build/transforms/RemoveTraitDefinitions.java

@@ -0,0 +1,82 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.build.transforms;
+
+import java.util.Collections;
+import java.util.Set;
+import java.util.function.Predicate;
+import software.amazon.smithy.build.TransformContext;
+import software.amazon.smithy.model.Model;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.transform.ModelTransformer;
+
+/**
+ * {@code removeTraitShapes} removes all trait definitions from a model,
+ * but leaves if the trait definition contains any of the provided
+ * {@code tags}.
+ *
+ * <p>This transformer will not remove prelude trait definitions.
+ */
+public final class RemoveTraitDefinitions extends ConfigurableProjectionTransformer<RemoveTraitDefinitions.Config> {
+
+    /**
+     * {@code removeTraitShapes} configuration settings.
+     */
+    public static final class Config {
+
+        private Set<String> exportTagged = Collections.emptySet();
+
+        /**
+         * You can <em>export</em> trait definitions by applying specific tags
+         * to the trait definition and adding the list of export tags as an
+         * argument to the transformer.
+         *
+         * @param exportByTags Tags that cause trait definitions to be
+         *                     exported.
+         */
+        public void setExportTagged(Set<String> exportByTags) {
+            this.exportTagged = exportByTags;
+        }
+
+        /**
+         * Gets the set of tags that are used to export trait definitions.
+         *
+         * @return the tags that are used to export trait definitions.
+         */
+        public Set<String> getExportTagged() {
+            return exportTagged;
+        }
+    }
+
+    @Override
+    public Class<Config> getConfigType() {
+        return Config.class;
+    }
+
+    @Override
+    protected Model transformWithConfig(TransformContext context, Config config) {
+        Model model = context.getModel();
+        ModelTransformer transformer = context.getTransformer();
+        Predicate<Shape> keepTraitDefsByTag = trait -> config.getExportTagged().stream().noneMatch(trait::hasTag);
+
+        return transformer.getModelWithoutTraitShapes(model, keepTraitDefsByTag);
+    }
+
+    @Override
+    public String getName() {
+        return "removeTraitDefinitions";
+    }
+}

smithy-build/src/test/java/software/amazon/smithy/build/transforms/RemoveTraitDefinitionsTest.java

@@ -0,0 +1,77 @@
+/*
+ * Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License").
+ * You may not use this file except in compliance with the License.
+ * A copy of the License is located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ * or in the "license" file accompanying this file. This file is distributed
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ * express or implied. See the License for the specific language governing
+ * permissions and limitations under the License.
+ */
+
+package software.amazon.smithy.build.transforms;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import java.nio.file.Paths;
+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.Shape;
+import software.amazon.smithy.model.shapes.ShapeId;
+
+public class RemoveTraitDefinitionsTest {
+
+    @Test
+    public void removesAllTraitDefinitions() throws Exception {
+        Model model = Model.assembler()
+                .addImport(Paths.get(getClass().getResource("remove-trait-definitions.json").toURI()))
+                .assemble()
+                .unwrap();
+        TransformContext context = TransformContext.builder()
+                .model(model)
+                .settings(Node.objectNode().withMember("exportTagged", Node.arrayNode()))
+                .build();
+        Model result = new RemoveTraitDefinitions().transform(context);
+        Shape baz = result.getShape(ShapeId.from("ns.foo#baz")).get();
+
+        // Trait definitions are removed.
+        assertFalse(result.getTraitDefinition(ShapeId.from("ns.foo#bar")).isPresent());
+        assertFalse(result.getShape(ShapeId.from("ns.foo#bar")).isPresent());
+        // Instance of trait on shape is retained.
+        assertTrue(baz.hasTrait("ns.foo#bar"));
+    }
+
+    @Test
+    public void retainsTaggedTraitDefinitions() throws Exception {
+        Model model = Model.assembler()
+                .addImport(Paths.get(getClass().getResource("retain-tagged-trait-definitions.json").toURI()))
+                .assemble()
+                .unwrap();
+        TransformContext context = TransformContext.builder()
+                .model(model)
+                .settings(Node.objectNode().withMember("exportTagged", Node.fromStrings("export")))
+                .build();
+        Model result = new RemoveTraitDefinitions().transform(context);
+        Shape baz = result.getShape(ShapeId.from("ns.foo#baz")).get();
+        Shape garply = result.getShape(ShapeId.from("ns.foo#garply")).get();
+
+        // Definition for trait "export" tag should be retained.
+        assertTrue(result.getTraitDefinition(ShapeId.from("ns.foo#bar")).isPresent());
+        assertTrue(result.getShape(ShapeId.from("ns.foo#bar")).isPresent());
+        // Instance of trait on shape is retained.
+        assertTrue(baz.hasTrait("ns.foo#bar"));
+
+        // Definition for trait without "export" tag should be removed.
+        assertFalse(result.getTraitDefinition(ShapeId.from("ns.foo#qux")).isPresent());
+        assertFalse(result.getShape(ShapeId.from("ns.foo#qux")).isPresent());
+        // Instance of trait on shape is retained.
+        assertTrue(garply.hasTrait("ns.foo#qux"));
+    }
+}

smithy-build/src/test/resources/software/amazon/smithy/build/transforms/remove-trait-definitions.json

@@ -0,0 +1,27 @@
+{
+  "smithy": "1.0",
+  "shapes": {
+    "ns.foo#bar": {
+      "type": "structure",
+      "members": {
+        "member": {
+          "target": "ns.foo#BarTraitShapeMember"
+        }
+      },
+      "traits": {
+        "smithy.api#trait": {}
+      }
+    },
+    "ns.foo#baz": {
+      "type": "string",
+      "traits": {
+        "ns.foo#bar": {
+          "member": "baz"
+        }
+      }
+    },
+    "ns.foo#BarTraitShapeMember": {
+      "type": "string"
+    }
+  }
+}

smithy-build/src/test/resources/software/amazon/smithy/build/transforms/retain-tagged-trait-definitions.json

@@ -0,0 +1,52 @@
+{
+  "smithy": "1.0",
+  "shapes": {
+    "ns.foo#bar": {
+      "type": "structure",
+      "members": {
+        "member": {
+          "target": "ns.foo#BarTraitShapeMember"
+        }
+      },
+      "traits": {
+        "smithy.api#trait": {},
+        "smithy.api#tags": [
+          "export"
+        ]
+      }
+    },
+    "ns.foo#baz": {
+      "type": "string",
+      "traits": {
+        "ns.foo#bar": {
+          "member": "baz"
+        }
+      }
+    },
+    "ns.foo#BarTraitShapeMember": {
+      "type": "string"
+    },
+    "ns.foo#qux": {
+      "type": "structure",
+      "members": {
+        "member": {
+          "target": "ns.foo#BarTraitShapeMember"
+        }
+      },
+      "traits": {
+        "smithy.api#trait": {},
+        "smithy.api#tags": [
+          "corge"
+        ]
+      }
+    },
+    "ns.foo#garply": {
+      "type": "string",
+      "traits": {
+        "ns.foo#qux": {
+          "member": "baz"
+        }
+      }
+    }
+  }
+}

smithy-model/src/main/java/software/amazon/smithy/model/transform/ModelTransformer.java

@@ -370,13 +370,32 @@ public Model removeUnreferencedTraitDefinitions(Model model, Predicate<Shape> ke
      *
      * <p>This can be useful when serializing a Smithy model to a format that
      * does not include trait definitions and the shapes used by trait definitions
-     * would have no meaning (e.g., Swagger).
+     * would have no meaning (e.g., OpenAPI).
      *
      * @param model Model to transform.
      * @return Returns the transformed model.base.
      */
     public Model scrubTraitDefinitions(Model model) {
-        return new ScrubTraitDefinitions().transform(this, model);
+        return scrubTraitDefinitions(model, FunctionalUtils.alwaysTrue());
+    }
+
+    /**
+     * Removes trait definitions from a model and all shapes that are
+     * only connected to the graph either directly or transitively by a
+     * trait definition shape.
+     *
+     * <p>This can be useful when serializing a Smithy model to a format that
+     * does not include trait definitions and the shapes used by trait definitions
+     * would have no meaning (e.g., OpenAPI).
+     *
+     * @param model Model to transform.
+     * @param keepFilter Predicate function that accepts an trait shape (that
+     *  has the {@link TraitDefinition} trait) and returns true to remove the
+     *  definition or false to keep the definition in the model.
+     * @return Returns the transformed model.
+     */
+    public Model scrubTraitDefinitions(Model model, Predicate<Shape> keepFilter) {
+        return new ScrubTraitDefinitions().transform(this, model, keepFilter);
     }
 
     /**
@@ -387,6 +406,20 @@ public Model scrubTraitDefinitions(Model model) {
      * @return Returns a model that contains matching shapes.
      */
     public Model getModelWithoutTraitShapes(Model model) {
+        return getModelWithoutTraitShapes(model, FunctionalUtils.alwaysTrue());
+    }
+
+    /**
+     * Gets all shapes from a model where shapes that define traits or shapes
+     * that are only used as part of a trait definition have been removed.
+     *
+     * @param model Model that contains shapes.
+     * @param keepFilter Predicate function that accepts a trait shape (that
+     *  has the {@link TraitDefinition} trait) and returns true to remove the
+     *  definition or false to keep the definition in the model.
+     * @return Returns a model that contains matching shapes.
+     */
+    public Model getModelWithoutTraitShapes(Model model, Predicate<Shape> keepFilter) {
         Model.Builder builder = Model.builder();
 
         // ScrubTraitDefinitions is used to removed traits and trait shapes.
@@ -395,7 +428,7 @@ public Model getModelWithoutTraitShapes(Model model) {
         // a model is created by getting all shape IDs from the modified
         // model, grabbing shapes from the original model, and building a new
         // Model.
-        scrubTraitDefinitions(model).shapes()
+        scrubTraitDefinitions(model, keepFilter).shapes()
                 .map(Shape::getId)
                 .map(model::getShape)
                 .map(Optional::get)