Skip to content

Commit

Permalink
Add CleanTraitDefinitions plugin (smithy-lang#379)
Browse files Browse the repository at this point in the history
  • Loading branch information
Chase Coalwell authored and JordonPhillips committed Apr 22, 2020
1 parent 3485a2e commit 07b1bd5
Show file tree
Hide file tree
Showing 4 changed files with 185 additions and 0 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,94 @@
/*
* 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.model.transform.plugins;

import java.util.Collection;
import java.util.List;
import java.util.Set;
import java.util.stream.Collectors;
import java.util.stream.Stream;
import software.amazon.smithy.model.Model;
import software.amazon.smithy.model.shapes.Shape;
import software.amazon.smithy.model.shapes.ShapeId;
import software.amazon.smithy.model.shapes.StructureShape;
import software.amazon.smithy.model.traits.AuthDefinitionTrait;
import software.amazon.smithy.model.traits.ProtocolDefinitionTrait;
import software.amazon.smithy.model.traits.Trait;
import software.amazon.smithy.model.transform.ModelTransformer;
import software.amazon.smithy.model.transform.ModelTransformerPlugin;

/**
* Removes traits from {@link AuthDefinitionTrait} and
* {@link ProtocolDefinitionTrait} traits that refer to removed shapes.
*/
public final class CleanTraitDefinitions implements ModelTransformerPlugin {
@Override
public Model onRemove(ModelTransformer transformer, Collection<Shape> removed, Model model) {
Set<ShapeId> removedShapeIds = removed.stream().map(Shape::getId).collect(Collectors.toSet());
model = transformer.replaceShapes(model, getAuthDefShapesToReplace(model, removedShapeIds));

return transformer.replaceShapes(model, getProtocolDefShapesToReplace(model, removedShapeIds));
}

private Set<Shape> getAuthDefShapesToReplace(Model model, Set<ShapeId> removedShapeIds) {
return model.shapes(StructureShape.class)
.flatMap(s -> Trait.flatMapStream(s, AuthDefinitionTrait.class))
.flatMap(pair -> {
AuthDefinitionTrait authDefTrait = pair.getRight();
List<ShapeId> traits = authDefTrait.getTraits();
List<ShapeId> newTraits = excludeTraitsInSet(traits, removedShapeIds);

// Return early if re-built list of traits is the same as existing list.
if (traits.equals(newTraits)) {
return Stream.empty();
}

// If the list of traits on the AuthDefinitionTrait has changed due to a trait shape being
// removed from the model, return a new version of the shape with a new version of the trait.
return Stream.of(pair.getLeft().toBuilder().addTrait(authDefTrait.toBuilder()
.traits(newTraits).build()).build());
})
.collect(Collectors.toSet());
}

private Set<Shape> getProtocolDefShapesToReplace(Model model, Set<ShapeId> removedShapeIds) {
return model.shapes(StructureShape.class)
.flatMap(s -> Trait.flatMapStream(s, ProtocolDefinitionTrait.class))
.flatMap(pair -> {
ProtocolDefinitionTrait protocolDefinitionTrait = pair.getRight();
List<ShapeId> traits = protocolDefinitionTrait.getTraits();
List<ShapeId> newTraits = excludeTraitsInSet(traits, removedShapeIds);

// Return early if re-built list of traits is the same as existing list.
if (traits.equals(newTraits)) {
return Stream.empty();
}

// If the list of traits on the ProtocolDefinitionTrait has changed due to a trait shape
// being removed from the model, return a new version of the shape with a new version of
// the trait.
return Stream.of(pair.getLeft().toBuilder().addTrait(protocolDefinitionTrait.toBuilder()
.traits(newTraits).build()).build());
})
.collect(Collectors.toSet());
}

private List<ShapeId> excludeTraitsInSet(List<ShapeId> traits, Set<ShapeId> shapeIds) {
return traits.stream()
.filter(trait -> !shapeIds.contains(trait))
.collect(Collectors.toList());
}
}
Original file line number Diff line number Diff line change
Expand Up @@ -2,4 +2,5 @@ software.amazon.smithy.model.transform.plugins.CleanBindings
software.amazon.smithy.model.transform.plugins.CleanOperationStructures
software.amazon.smithy.model.transform.plugins.CleanResourceReferences
software.amazon.smithy.model.transform.plugins.CleanStructureAndUnionMembers
software.amazon.smithy.model.transform.plugins.CleanTraitDefinitions
software.amazon.smithy.model.transform.plugins.RemoveTraits
Original file line number Diff line number Diff line change
Expand Up @@ -16,6 +16,7 @@
package software.amazon.smithy.model.transform;

import static org.hamcrest.MatcherAssert.assertThat;
import static org.junit.jupiter.api.Assertions.assertFalse;

import java.util.Arrays;
import java.util.Collections;
Expand All @@ -35,6 +36,8 @@
import software.amazon.smithy.model.shapes.StringShape;
import software.amazon.smithy.model.shapes.StructureShape;
import software.amazon.smithy.model.shapes.UnionShape;
import software.amazon.smithy.model.traits.AuthDefinitionTrait;
import software.amazon.smithy.model.traits.ProtocolDefinitionTrait;
import software.amazon.smithy.model.traits.ReadonlyTrait;

public class RemoveShapesTest {
Expand Down Expand Up @@ -140,4 +143,44 @@ public void removesOperationsFromResourcesWhenOperationRemoved() {
assertThat(result.expectShape(container.getId()).asResourceShape().get().getOperations(),
Matchers.contains(c.getId()));
}

@Test
public void removesTraitsFromAuthDefinitionWhenReferenceRemoved() {
Model model = Model.assembler()
.addImport(getClass().getResource("remove-shapes.json"))
.assemble()
.unwrap();
ShapeId removedId = ShapeId.from("ns.foo#bar");
Shape removedShape = model.expectShape(removedId);


ModelTransformer transformer = ModelTransformer.create();
Model result = transformer.removeShapes(model, Collections.singletonList(removedShape));

ShapeId subjectId = ShapeId.from("ns.foo#auth");
Shape subject = result.expectShape(subjectId);
AuthDefinitionTrait trait = subject.getTrait(AuthDefinitionTrait.class).get();

assertFalse(trait.getTraits().contains(removedId));
}

@Test
public void removesTraitsFromProtocolDefinitionWhenReferenceRemoved() {
Model model = Model.assembler()
.addImport(getClass().getResource("remove-shapes.json"))
.assemble()
.unwrap();
ShapeId removedId = ShapeId.from("ns.foo#baz");
Shape removedShape = model.expectShape(removedId);


ModelTransformer transformer = ModelTransformer.create();
Model result = transformer.removeShapes(model, Collections.singletonList(removedShape));

ShapeId subjectId = ShapeId.from("ns.foo#protocol");
Shape subject = result.expectShape(subjectId);
ProtocolDefinitionTrait trait = subject.getTrait(ProtocolDefinitionTrait.class).get();

assertFalse(trait.getTraits().contains(removedId));
}
}
Original file line number Diff line number Diff line change
@@ -0,0 +1,47 @@
{
"smithy": "1.0.0",
"shapes": {
"ns.foo#auth": {
"type": "structure",
"traits": {
"smithy.api#trait": {
"selector": "service"
},
"smithy.api#authDefinition": {
"traits": [
"ns.foo#bar"
]
}
}
},
"ns.foo#bar": {
"type": "structure",
"traits": {
"smithy.api#trait": {
"selector": "operation"
}
}
},
"ns.foo#protocol": {
"type": "structure",
"traits": {
"smithy.api#trait": {
"selector": "service"
},
"smithy.api#protocolDefinition": {
"traits": [
"ns.foo#baz"
]
}
}
},
"ns.foo#baz": {
"type": "structure",
"traits": {
"smithy.api#trait": {
"selector": "operation"
}
}
}
}
}

0 comments on commit 07b1bd5

Please sign in to comment.