Skip to content

Commit 82c1730

Browse files
author
Chase Coalwell
authored
Add CleanTraitDefinitions plugin (#379)
1 parent c40a385 commit 82c1730

File tree

4 files changed

+185
-0
lines changed

4 files changed

+185
-0
lines changed

‎smithy-model/src/main/java/software/amazon/smithy/model/transform/plugins/CleanTraitDefinitions.java

+94
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,94 @@
1+
/*
2+
* Copyright 2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
3+
*
4+
* Licensed under the Apache License, Version 2.0 (the "License").
5+
* You may not use this file except in compliance with the License.
6+
* A copy of the License is located at
7+
*
8+
* http://aws.amazon.com/apache2.0
9+
*
10+
* or in the "license" file accompanying this file. This file is distributed
11+
* on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12+
* express or implied. See the License for the specific language governing
13+
* permissions and limitations under the License.
14+
*/
15+
16+
package software.amazon.smithy.model.transform.plugins;
17+
18+
import java.util.Collection;
19+
import java.util.List;
20+
import java.util.Set;
21+
import java.util.stream.Collectors;
22+
import java.util.stream.Stream;
23+
import software.amazon.smithy.model.Model;
24+
import software.amazon.smithy.model.shapes.Shape;
25+
import software.amazon.smithy.model.shapes.ShapeId;
26+
import software.amazon.smithy.model.shapes.StructureShape;
27+
import software.amazon.smithy.model.traits.AuthDefinitionTrait;
28+
import software.amazon.smithy.model.traits.ProtocolDefinitionTrait;
29+
import software.amazon.smithy.model.traits.Trait;
30+
import software.amazon.smithy.model.transform.ModelTransformer;
31+
import software.amazon.smithy.model.transform.ModelTransformerPlugin;
32+
33+
/**
34+
* Removes traits from {@link AuthDefinitionTrait} and
35+
* {@link ProtocolDefinitionTrait} traits that refer to removed shapes.
36+
*/
37+
public final class CleanTraitDefinitions implements ModelTransformerPlugin {
38+
@Override
39+
public Model onRemove(ModelTransformer transformer, Collection<Shape> removed, Model model) {
40+
Set<ShapeId> removedShapeIds = removed.stream().map(Shape::getId).collect(Collectors.toSet());
41+
model = transformer.replaceShapes(model, getAuthDefShapesToReplace(model, removedShapeIds));
42+
43+
return transformer.replaceShapes(model, getProtocolDefShapesToReplace(model, removedShapeIds));
44+
}
45+
46+
private Set<Shape> getAuthDefShapesToReplace(Model model, Set<ShapeId> removedShapeIds) {
47+
return model.shapes(StructureShape.class)
48+
.flatMap(s -> Trait.flatMapStream(s, AuthDefinitionTrait.class))
49+
.flatMap(pair -> {
50+
AuthDefinitionTrait authDefTrait = pair.getRight();
51+
List<ShapeId> traits = authDefTrait.getTraits();
52+
List<ShapeId> newTraits = excludeTraitsInSet(traits, removedShapeIds);
53+
54+
// Return early if re-built list of traits is the same as existing list.
55+
if (traits.equals(newTraits)) {
56+
return Stream.empty();
57+
}
58+
59+
// If the list of traits on the AuthDefinitionTrait has changed due to a trait shape being
60+
// removed from the model, return a new version of the shape with a new version of the trait.
61+
return Stream.of(pair.getLeft().toBuilder().addTrait(authDefTrait.toBuilder()
62+
.traits(newTraits).build()).build());
63+
})
64+
.collect(Collectors.toSet());
65+
}
66+
67+
private Set<Shape> getProtocolDefShapesToReplace(Model model, Set<ShapeId> removedShapeIds) {
68+
return model.shapes(StructureShape.class)
69+
.flatMap(s -> Trait.flatMapStream(s, ProtocolDefinitionTrait.class))
70+
.flatMap(pair -> {
71+
ProtocolDefinitionTrait protocolDefinitionTrait = pair.getRight();
72+
List<ShapeId> traits = protocolDefinitionTrait.getTraits();
73+
List<ShapeId> newTraits = excludeTraitsInSet(traits, removedShapeIds);
74+
75+
// Return early if re-built list of traits is the same as existing list.
76+
if (traits.equals(newTraits)) {
77+
return Stream.empty();
78+
}
79+
80+
// If the list of traits on the ProtocolDefinitionTrait has changed due to a trait shape
81+
// being removed from the model, return a new version of the shape with a new version of
82+
// the trait.
83+
return Stream.of(pair.getLeft().toBuilder().addTrait(protocolDefinitionTrait.toBuilder()
84+
.traits(newTraits).build()).build());
85+
})
86+
.collect(Collectors.toSet());
87+
}
88+
89+
private List<ShapeId> excludeTraitsInSet(List<ShapeId> traits, Set<ShapeId> shapeIds) {
90+
return traits.stream()
91+
.filter(trait -> !shapeIds.contains(trait))
92+
.collect(Collectors.toList());
93+
}
94+
}

‎smithy-model/src/main/resources/META-INF/services/software.amazon.smithy.model.transform.ModelTransformerPlugin

+1
Original file line numberDiff line numberDiff line change
@@ -2,4 +2,5 @@ software.amazon.smithy.model.transform.plugins.CleanBindings
22
software.amazon.smithy.model.transform.plugins.CleanOperationStructures
33
software.amazon.smithy.model.transform.plugins.CleanResourceReferences
44
software.amazon.smithy.model.transform.plugins.CleanStructureAndUnionMembers
5+
software.amazon.smithy.model.transform.plugins.CleanTraitDefinitions
56
software.amazon.smithy.model.transform.plugins.RemoveTraits

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

+43
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
package software.amazon.smithy.model.transform;
1717

1818
import static org.hamcrest.MatcherAssert.assertThat;
19+
import static org.junit.jupiter.api.Assertions.assertFalse;
1920

2021
import java.util.Arrays;
2122
import java.util.Collections;
@@ -35,6 +36,8 @@
3536
import software.amazon.smithy.model.shapes.StringShape;
3637
import software.amazon.smithy.model.shapes.StructureShape;
3738
import software.amazon.smithy.model.shapes.UnionShape;
39+
import software.amazon.smithy.model.traits.AuthDefinitionTrait;
40+
import software.amazon.smithy.model.traits.ProtocolDefinitionTrait;
3841
import software.amazon.smithy.model.traits.ReadonlyTrait;
3942

4043
public class RemoveShapesTest {
@@ -140,4 +143,44 @@ public void removesOperationsFromResourcesWhenOperationRemoved() {
140143
assertThat(result.expectShape(container.getId()).asResourceShape().get().getOperations(),
141144
Matchers.contains(c.getId()));
142145
}
146+
147+
@Test
148+
public void removesTraitsFromAuthDefinitionWhenReferenceRemoved() {
149+
Model model = Model.assembler()
150+
.addImport(getClass().getResource("remove-shapes.json"))
151+
.assemble()
152+
.unwrap();
153+
ShapeId removedId = ShapeId.from("ns.foo#bar");
154+
Shape removedShape = model.expectShape(removedId);
155+
156+
157+
ModelTransformer transformer = ModelTransformer.create();
158+
Model result = transformer.removeShapes(model, Collections.singletonList(removedShape));
159+
160+
ShapeId subjectId = ShapeId.from("ns.foo#auth");
161+
Shape subject = result.expectShape(subjectId);
162+
AuthDefinitionTrait trait = subject.getTrait(AuthDefinitionTrait.class).get();
163+
164+
assertFalse(trait.getTraits().contains(removedId));
165+
}
166+
167+
@Test
168+
public void removesTraitsFromProtocolDefinitionWhenReferenceRemoved() {
169+
Model model = Model.assembler()
170+
.addImport(getClass().getResource("remove-shapes.json"))
171+
.assemble()
172+
.unwrap();
173+
ShapeId removedId = ShapeId.from("ns.foo#baz");
174+
Shape removedShape = model.expectShape(removedId);
175+
176+
177+
ModelTransformer transformer = ModelTransformer.create();
178+
Model result = transformer.removeShapes(model, Collections.singletonList(removedShape));
179+
180+
ShapeId subjectId = ShapeId.from("ns.foo#protocol");
181+
Shape subject = result.expectShape(subjectId);
182+
ProtocolDefinitionTrait trait = subject.getTrait(ProtocolDefinitionTrait.class).get();
183+
184+
assertFalse(trait.getTraits().contains(removedId));
185+
}
143186
}

‎smithy-model/src/test/resources/software/amazon/smithy/model/transform/remove-shapes.json

+47
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
{
2+
"smithy": "1.0.0",
3+
"shapes": {
4+
"ns.foo#auth": {
5+
"type": "structure",
6+
"traits": {
7+
"smithy.api#trait": {
8+
"selector": "service"
9+
},
10+
"smithy.api#authDefinition": {
11+
"traits": [
12+
"ns.foo#bar"
13+
]
14+
}
15+
}
16+
},
17+
"ns.foo#bar": {
18+
"type": "structure",
19+
"traits": {
20+
"smithy.api#trait": {
21+
"selector": "operation"
22+
}
23+
}
24+
},
25+
"ns.foo#protocol": {
26+
"type": "structure",
27+
"traits": {
28+
"smithy.api#trait": {
29+
"selector": "service"
30+
},
31+
"smithy.api#protocolDefinition": {
32+
"traits": [
33+
"ns.foo#baz"
34+
]
35+
}
36+
}
37+
},
38+
"ns.foo#baz": {
39+
"type": "structure",
40+
"traits": {
41+
"smithy.api#trait": {
42+
"selector": "operation"
43+
}
44+
}
45+
}
46+
}
47+
}

0 commit comments

Comments
 (0)