Improve member trait querys and fix JSON bug

smithy-jsonschema/src/main/java/software/amazon/smithy/jsonschema/JsonSchemaConverter.java

@@ -27,6 +27,7 @@
 import software.amazon.smithy.model.shapes.ServiceShape;
 import software.amazon.smithy.model.shapes.Shape;
 import software.amazon.smithy.model.shapes.ShapeIndex;
+import software.amazon.smithy.model.traits.EffectiveTraitQuery;
 import software.amazon.smithy.model.traits.PrivateTrait;
 import software.amazon.smithy.utils.FunctionalUtils;
 import software.amazon.smithy.utils.Pair;
@@ -267,15 +268,14 @@ private boolean isExcludedPrivateShape(ShapeIndex shapeIndex, Shape shape) {
         // We can explicitly enable the generation of private shapes if desired.
         if (getConfig().getBooleanMemberOrDefault(JsonSchemaConstants.SMITHY_INCLUDE_PRIVATE_SHAPES)) {
             return false;
-        } else if (shape.hasTrait(PrivateTrait.class)) {
-            return true;
         }
 
-        // Now check members.
-        return shape.asMemberShape()
-                .flatMap(member -> shapeIndex.getShape(member.getContainer()))
-                .filter(parent -> parent.hasTrait(PrivateTrait.class))
-                .isPresent();
+        return EffectiveTraitQuery.builder()
+                .shapeIndex(shapeIndex)
+                .traitClass(PrivateTrait.class)
+                .inheritFromContainer(true)
+                .build()
+                .isTraitApplied(shape);
     }
 
     private void addExtensions(SchemaDocument.Builder builder) {

smithy-jsonschema/src/test/java/software/amazon/smithy/jsonschema/JsonSchemaConverterTest.java

@@ -198,6 +198,18 @@ public void excludesMembersOfPrivateShapes() {
         assertThat(doc.getDefinitions().keySet(), contains("#/definitions/SmithyExampleString"));
     }
 
+    @Test
+    public void excludesMembersThatTargetPrivateShapes() {
+        StringShape string = StringShape.builder().id("smithy.example#String").addTrait(new PrivateTrait()).build();
+        MemberShape member = MemberShape.builder().id("smithy.example#Foo$bar").target(string).build();
+        StructureShape struct = StructureShape.builder().id("smithy.example#Foo").addMember(member).build();
+        ShapeIndex index = ShapeIndex.builder().addShapes(struct, member, string).build();
+        SchemaDocument doc = JsonSchemaConverter.create().convert(index);
+
+        // The member and the target are filtered out.
+        assertThat(doc.getDefinitions().keySet(), contains("#/definitions/SmithyExampleFoo"));
+    }
+
     @Test
     public void canIncludePrivateShapesWithFlag() {
         StringShape string = StringShape.builder().id("smithy.example#String").build();

smithy-model/src/main/java/software/amazon/smithy/model/knowledge/HttpBindingIndex.java

@@ -236,7 +236,6 @@ public TimestampFormatTrait.Format determineTimestampFormat(
             TimestampFormatTrait.Format defaultFormat
     ) {
         return index.getShape(member.toShapeId())
-                .flatMap(Shape::asMemberShape)
                 // Use the timestampFormat trait on the member or target if present.
                 .flatMap(shape -> shape.getMemberTrait(index, TimestampFormatTrait.class))
                 .map(TimestampFormatTrait::getFormat)

smithy-model/src/main/java/software/amazon/smithy/model/shapes/MemberShape.java

@@ -100,28 +100,13 @@ public boolean equals(Object other) {
         return super.equals(other) && getTarget().equals(((MemberShape) other).getTarget());
     }
 
-    /**
-     * Gets a trait from the member shape or from the shape targeted by the
-     * member.
-     *
-     * @param index Shape index used to find member targets.
-     * @param trait Trait type to get.
-     * @param <T> Trait type to get.
-     * @return Returns the optionally found trait on the shape or member.
-     */
+    @Override
     public <T extends Trait> Optional<T> getMemberTrait(ShapeIndex index, Class<T> trait) {
         return OptionalUtils.or(getTrait(trait),
                 () -> index.getShape(getTarget()).flatMap(targetedShape -> targetedShape.getTrait(trait)));
     }
 
-    /**
-     * Gets a trait from the member shape or from the shape targeted by the
-     * member.
-     *
-     * @param index Shape index used to find member targets.
-     * @param traitName Trait name to get.
-     * @return Returns the optionally found trait on the shape or member.
-     */
+    @Override
     public Optional<Trait> findMemberTrait(ShapeIndex index, String traitName) {
         return OptionalUtils.or(findTrait(traitName),
                 () -> index.getShape(getTarget()).flatMap(targetedShape -> targetedShape.findTrait(traitName)));

smithy-model/src/main/java/software/amazon/smithy/model/shapes/Shape.java

@@ -214,6 +214,39 @@ public final Map<ShapeId, Trait> getAllTraits() {
         return traits;
     }
 
+    /**
+     * Gets a trait from the member shape or from the shape targeted by the
+     * member.
+     *
+     * <p>If the shape is not a member, then the method functions the same as
+     * {@link #getTrait(Class)}.
+     *
+     * @param index Shape index used to find member targets.
+     * @param trait Trait type to get.
+     * @param <T> Trait type to get.
+     * @return Returns the optionally found trait on the shape or member.
+     * @see MemberShape#getTrait(Class)
+     */
+    public <T extends Trait> Optional<T> getMemberTrait(ShapeIndex index, Class<T> trait) {
+        return getTrait(trait);
+    }
+
+    /**
+     * Gets a trait from the member shape or from the shape targeted by the
+     * member.
+     *
+     * <p>If the shape is not a member, then the method functions the same as
+     * {@link #findTrait(String)}.
+     *
+     * @param index Shape index used to find member targets.
+     * @param traitName Trait name to get.
+     * @return Returns the optionally found trait on the shape or member.
+     * @see MemberShape#findTrait(String)
+     */
+    public Optional<Trait> findMemberTrait(ShapeIndex index, String traitName) {
+        return findTrait(traitName);
+    }
+
     /**
      * @return Optionally returns the shape as a {@link BigDecimalShape}.
      */

smithy-model/src/main/java/software/amazon/smithy/model/traits/EffectiveTraitQuery.java

@@ -0,0 +1,134 @@
+/*
+ * Copyright 2019 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.model.traits;
+
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeIndex;
+import software.amazon.smithy.model.shapes.ToShapeId;
+import software.amazon.smithy.utils.SmithyBuilder;
+import software.amazon.smithy.utils.ToSmithyBuilder;
+
+/**
+ * Queries a shape index for effective traits bound to shapes and members.
+ */
+public final class EffectiveTraitQuery implements ToSmithyBuilder<EffectiveTraitQuery> {
+
+    private final ShapeIndex shapeIndex;
+    private final Class<? extends Trait> traitClass;
+    private final boolean inheritFromContainer;
+
+    private EffectiveTraitQuery(Builder builder) {
+        this.shapeIndex = SmithyBuilder.requiredState("shapeIndex", builder.shapeIndex);
+        this.traitClass = SmithyBuilder.requiredState("traitClass", builder.traitClass);
+        this.inheritFromContainer = builder.inheritFromContainer;
+    }
+
+    /**
+     * Checks if the trait is effectively applied to a shape.
+     *
+     * @param shapeId Shape to test.
+     * @return Returns true if the trait is effectively applied to the shape.
+     */
+    public boolean isTraitApplied(ToShapeId shapeId) {
+        Shape shape = shapeIndex.getShape(shapeId.toShapeId()).orElse(null);
+
+        if (shape == null) {
+            return false;
+        }
+
+        if (shape.getMemberTrait(shapeIndex, traitClass).isPresent()) {
+            return true;
+        }
+
+        if (!inheritFromContainer || !shape.asMemberShape().isPresent()) {
+            return false;
+        }
+
+        // Check if the parent of the member is marked with the trait.
+        MemberShape memberShape = shape.asMemberShape().get();
+        Shape parent = shapeIndex.getShape(memberShape.getContainer()).orElse(null);
+        return parent != null && parent.hasTrait(traitClass);
+    }
+
+    /**
+     * Creates a new query builder.
+     *
+     * @return Returns the created builder.
+     */
+    public static Builder builder() {
+        return new Builder();
+    }
+
+    @Override
+    public Builder toBuilder() {
+        return builder()
+                .shapeIndex(shapeIndex)
+                .traitClass(traitClass)
+                .inheritFromContainer(inheritFromContainer);
+    }
+
+    /**
+     * Builds a reusable EffectiveTraitQuery.
+     */
+    public static final class Builder implements SmithyBuilder<EffectiveTraitQuery> {
+
+        private ShapeIndex shapeIndex;
+        private Class<? extends Trait> traitClass;
+        private boolean inheritFromContainer;
+
+        @Override
+        public EffectiveTraitQuery build() {
+            return new EffectiveTraitQuery(this);
+        }
+
+        /**
+         * Sets the required shape index to query.
+         *
+         * @param shapeIndex Shape index to query.
+         * @return Returns the query object builder.
+         */
+        public Builder shapeIndex(ShapeIndex shapeIndex) {
+            this.shapeIndex = shapeIndex;
+            return this;
+        }
+
+        /**
+         * Sets the required trait being queried.
+         *
+         * @param traitClass Trait to detect on shapes.
+         * @return Returns the query object builder.
+         */
+        public Builder traitClass(Class<? extends Trait> traitClass) {
+            this.traitClass = traitClass;
+            return this;
+        }
+
+        /**
+         * When testing member shapes, also checks the container of the member for
+         * the presence of a trait.
+         *
+         * <p>By default, traits are not inherited from a member's parent container.
+         *
+         * @param inheritFromContainer Set to true to inherit traits from member containers.
+         * @return Returns the query object builder.
+         */
+        public Builder inheritFromContainer(boolean inheritFromContainer) {
+            this.inheritFromContainer = inheritFromContainer;
+            return this;
+        }
+    }
+}

smithy-model/src/test/java/software/amazon/smithy/model/shapes/ShapeTest.java

@@ -145,11 +145,19 @@ public void hasTraits() {
                 .addTrait(otherTrait)
                 .addTrait(documentationTrait)
                 .build();
+        ShapeIndex index = ShapeIndex.builder()
+                .addShapes(shape)
+                .build();
 
         assertTrue(shape.getTrait(MyTrait.class).isPresent());
+        assertTrue(shape.getMemberTrait(index, MyTrait.class).isPresent());
+
         assertTrue(shape.findTrait("foo.baz#foo").isPresent());
+        assertTrue(shape.findMemberTrait(index, "foo.baz#foo").isPresent());
+
         assertTrue(shape.hasTrait("foo.baz#foo"));
         assertTrue(shape.getTrait(OtherTrait.class).isPresent());
+
         assertFalse(shape.getTrait(AnotherTrait.class).isPresent());
         assertFalse(shape.findTrait("notThere").isPresent());
         assertFalse(shape.hasTrait("notThere"));

smithy-model/src/test/java/software/amazon/smithy/model/traits/EffectiveTraitQueryTest.java

@@ -0,0 +1,99 @@
+package software.amazon.smithy.model.traits;
+
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertTrue;
+
+import org.junit.jupiter.api.Test;
+import software.amazon.smithy.model.shapes.ListShape;
+import software.amazon.smithy.model.shapes.MemberShape;
+import software.amazon.smithy.model.shapes.Shape;
+import software.amazon.smithy.model.shapes.ShapeIndex;
+import software.amazon.smithy.model.shapes.StringShape;
+
+public class EffectiveTraitQueryTest {
+    @Test
+    public void detectsTraitOnShape() {
+        Shape stringShape = StringShape.builder().id("foo.bar#Baz")
+                .addTrait(new SensitiveTrait())
+                .build();
+        ShapeIndex index = ShapeIndex.builder()
+                .addShapes(stringShape)
+                .build();
+        EffectiveTraitQuery query = EffectiveTraitQuery.builder()
+                .shapeIndex(index)
+                .traitClass(SensitiveTrait.class)
+                .build();
+
+        assertTrue(query.isTraitApplied(stringShape));
+    }
+
+    @Test
+    public void detectsTraitOnMemberTarget() {
+        Shape stringShape = StringShape.builder().id("foo.bar#Baz")
+                .addTrait(new SensitiveTrait())
+                .build();
+        MemberShape member = MemberShape.builder()
+                .id("foo.baz#Container$member")
+                .target(stringShape)
+                .build();
+        ShapeIndex index = ShapeIndex.builder()
+                .addShapes(stringShape, member)
+                .build();
+        EffectiveTraitQuery query = EffectiveTraitQuery.builder()
+                .shapeIndex(index)
+                .traitClass(SensitiveTrait.class)
+                .build();
+
+        assertTrue(query.isTraitApplied(member));
+    }
+
+    @Test
+    public void ignoresTraitOnMemberContainerByDefault() {
+        Shape stringShape = StringShape.builder().id("foo.bar#Baz").build();
+        MemberShape member = MemberShape.builder()
+                .id("foo.baz#Container$member")
+                .target(stringShape)
+                .build();
+        ListShape list = ListShape.builder()
+                .id("foo.baz#Container")
+                .member(member)
+                .addTrait(new SensitiveTrait())
+                .build();
+        ShapeIndex index = ShapeIndex.builder()
+                .addShapes(stringShape, member, list)
+                .build();
+        EffectiveTraitQuery query = EffectiveTraitQuery.builder()
+                .shapeIndex(index)
+                .traitClass(SensitiveTrait.class)
+                .build();
+
+        assertFalse(query.isTraitApplied(member));
+    }
+
+    @Test
+    public void detectsTraitOnMemberContainer() {
+        Shape stringShape = StringShape.builder().id("foo.bar#Baz").build();
+        MemberShape member = MemberShape.builder()
+                .id("foo.baz#Container$member")
+                .target(stringShape)
+                .build();
+        ListShape list = ListShape.builder()
+                .id("foo.baz#Container")
+                .member(member)
+                .addTrait(new SensitiveTrait())
+                .build();
+        ShapeIndex index = ShapeIndex.builder()
+                .addShapes(stringShape, member, list)
+                .build();
+        EffectiveTraitQuery query = EffectiveTraitQuery.builder()
+                .shapeIndex(index)
+                .traitClass(SensitiveTrait.class)
+                .inheritFromContainer(true)
+                .build();
+
+        assertTrue(query.isTraitApplied(member));
+
+        // Converts to a builder...
+        assertTrue(query.toBuilder().build().isTraitApplied(member));
+    }
+}