Skip to content

Commit 39a51b7

Browse files
Dorenavantlandonxjamesysaito1001
authored
Add EnumSection to allow decorators to modify enum member attributes (#4039)
## Motivation and Context <!--- Why is this change required? What problem does it solve? --> Allows users to use decorators to add additional attributes to members of an enum. <!--- If it fixes an open issue, please link to the issue here --> N/A ## Description <!--- Describe your changes in detail --> Adds `EnumSection` with named `AdditionalMemberAttributes` to allow decorators to modify enum codegen. ## Testing <!--- Please describe in detail how you tested your changes --> Could not find existing unit tests for customizations; please advise if you would like them to be added here. Ran `./gradlew` to ensure current tests pass. <!--- Include details of your testing environment, and the tests you ran to --> <!--- see how your change affects other areas of the code, etc. --> ## Checklist <!--- If a checkbox below is not applicable, then please DELETE it rather than leaving it unchecked --> - [X] For changes to the smithy-rs codegen or runtime crates, I have created a changelog entry Markdown file in the `.changelog` directory, specifying "client," "server," or both in the `applies_to` key. ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ --------- Co-authored-by: Landon James <lnj@amazon.com> Co-authored-by: ysaito1001 <awsaito@amazon.com>
1 parent de4be56 commit 39a51b7

File tree

20 files changed

+149
-50
lines changed

20 files changed

+149
-50
lines changed

‎.changelog/1740703869.md‎

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,10 @@
1+
---
2+
applies_to: ["client", "server"]
3+
authors: [Dorenavant]
4+
references: []
5+
breaking: false
6+
new_feature: true
7+
bug_fix: false
8+
---
9+
10+
Add EnumSection to allow decorators to modify enum member attributes

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt‎

Lines changed: 22 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -93,7 +93,14 @@ class ClientCodegenVisitor(
9393
model = codegenDecorator.transformModel(untransformedService, baseModel, settings)
9494
// the model transformer _might_ change the service shape
9595
val service = settings.getService(model)
96-
symbolProvider = RustClientCodegenPlugin.baseSymbolProvider(settings, model, service, rustSymbolProviderConfig, codegenDecorator)
96+
symbolProvider =
97+
RustClientCodegenPlugin.baseSymbolProvider(
98+
settings,
99+
model,
100+
service,
101+
rustSymbolProviderConfig,
102+
codegenDecorator,
103+
)
97104

98105
codegenContext =
99106
ClientCodegenContext(
@@ -177,7 +184,10 @@ class ClientCodegenVisitor(
177184
)
178185
try {
179186
// use an increased max_width to make rustfmt fail less frequently
180-
"cargo fmt -- --config max_width=150".runCommand(fileManifest.baseDir, timeout = settings.codegenConfig.formatTimeoutSeconds.toLong())
187+
"cargo fmt -- --config max_width=150".runCommand(
188+
fileManifest.baseDir,
189+
timeout = settings.codegenConfig.formatTimeoutSeconds.toLong(),
190+
)
181191
} catch (err: CommandError) {
182192
logger.warning("Failed to run cargo fmt: [${service.id}]\n${err.output}")
183193
}
@@ -236,7 +246,10 @@ class ClientCodegenVisitor(
236246

237247
implBlock(symbolProvider.toSymbol(shape)) {
238248
BuilderGenerator.renderConvenienceMethod(this, symbolProvider, shape)
239-
if (codegenContext.protocolImpl?.httpBindingResolver?.handlesEventStreamInitialResponse(shape) == true) {
249+
if (codegenContext.protocolImpl?.httpBindingResolver?.handlesEventStreamInitialResponse(
250+
shape,
251+
) == true
252+
) {
240253
BuilderGenerator.renderIntoBuilderMethod(this, symbolProvider, shape)
241254
}
242255
}
@@ -251,6 +264,7 @@ class ClientCodegenVisitor(
251264
}
252265
struct to builder
253266
}
267+
254268
else -> {
255269
val errorGenerator =
256270
ErrorGenerator(
@@ -283,7 +297,11 @@ class ClientCodegenVisitor(
283297
if (shape.hasTrait<EnumTrait>()) {
284298
val privateModule = privateModule(shape)
285299
rustCrate.inPrivateModuleWithReexport(privateModule, symbolProvider.toSymbol(shape)) {
286-
ClientEnumGenerator(codegenContext, shape).render(this)
300+
ClientEnumGenerator(
301+
codegenContext,
302+
shape,
303+
codegenDecorator.enumCustomizations(codegenContext, emptyList()),
304+
).render(this)
287305
}
288306
}
289307
}

‎codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
1919
import software.amazon.smithy.rust.codegen.core.rustlang.writable
2020
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
2121
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
22+
import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumCustomization
2223
import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGenerator
2324
import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGeneratorContext
2425
import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumMemberModel
@@ -283,19 +284,24 @@ data class InfallibleEnumType(
283284
}
284285
}
285286

286-
class ClientEnumGenerator(codegenContext: ClientCodegenContext, shape: StringShape) :
287+
class ClientEnumGenerator(
288+
codegenContext: ClientCodegenContext,
289+
shape: StringShape,
290+
customizations: List<EnumCustomization>,
291+
) :
287292
EnumGenerator(
288-
codegenContext.model,
289-
codegenContext.symbolProvider,
290-
shape,
291-
InfallibleEnumType(
292-
RustModule.new(
293-
"sealed_enum_unknown",
294-
visibility = Visibility.PUBCRATE,
295-
parent = ClientRustModule.primitives,
293+
codegenContext.model,
294+
codegenContext.symbolProvider,
295+
shape,
296+
InfallibleEnumType(
297+
RustModule.new(
298+
"sealed_enum_unknown",
299+
visibility = Visibility.PUBCRATE,
300+
parent = ClientRustModule.primitives,
301+
),
296302
),
297-
),
298-
)
303+
customizations,
304+
)
299305

300306
private fun unknownVariantError(): RuntimeType =
301307
RuntimeType.forInlineFun("UnknownVariantError", ClientRustModule.Error) {

‎codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGeneratorTest.kt‎

Lines changed: 6 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -28,7 +28,7 @@ class ClientEnumGeneratorTest {
2828
val context = testClientCodegenContext(model)
2929
val project = TestWorkspace.testProject(context.symbolProvider)
3030
project.moduleFor(shape) {
31-
ClientEnumGenerator(context, shape).render(this)
31+
ClientEnumGenerator(context, shape, emptyList()).render(this)
3232
unitTest(
3333
"matching_on_enum_should_be_forward_compatible",
3434
"""
@@ -88,7 +88,7 @@ class ClientEnumGeneratorTest {
8888
val context = testClientCodegenContext(model)
8989
val project = TestWorkspace.testProject(context.symbolProvider)
9090
project.moduleFor(shape) {
91-
ClientEnumGenerator(context, shape).render(this)
91+
ClientEnumGenerator(context, shape, emptyList()).render(this)
9292
unitTest(
9393
"impl_debug_for_non_sensitive_enum_should_implement_the_derived_debug_trait",
9494
"""
@@ -134,10 +134,10 @@ class ClientEnumGeneratorTest {
134134
val context = testClientCodegenContext(model)
135135
val project = TestWorkspace.testProject(context.symbolProvider)
136136
project.moduleFor(shapeA) {
137-
ClientEnumGenerator(context, shapeA).render(this)
137+
ClientEnumGenerator(context, shapeA, emptyList()).render(this)
138138
}
139139
project.moduleFor(shapeB) {
140-
ClientEnumGenerator(context, shapeB).render(this)
140+
ClientEnumGenerator(context, shapeB, emptyList()).render(this)
141141
unitTest(
142142
"impl_debug_for_non_sensitive_enum_should_implement_the_derived_debug_trait",
143143
"""
@@ -172,7 +172,7 @@ class ClientEnumGeneratorTest {
172172
val context = testClientCodegenContext(model)
173173
val project = TestWorkspace.testProject(context.symbolProvider)
174174
project.moduleFor(shape) {
175-
ClientEnumGenerator(context, shape).render(this)
175+
ClientEnumGenerator(context, shape, emptyList()).render(this)
176176
unitTest(
177177
"it_escapes_the_unknown_variant_if_the_enum_has_an_unknown_value_in_the_model",
178178
"""
@@ -205,7 +205,7 @@ class ClientEnumGeneratorTest {
205205
val project = TestWorkspace.testProject(context.symbolProvider)
206206
project.moduleFor(shape) {
207207
rust("##![allow(deprecated)]")
208-
ClientEnumGenerator(context, shape).render(this)
208+
ClientEnumGenerator(context, shape, emptyList()).render(this)
209209
unitTest(
210210
"generated_named_enums_roundtrip",
211211
"""

‎codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientInstantiatorTest.kt‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -53,7 +53,7 @@ internal class ClientInstantiatorTest {
5353

5454
val project = TestWorkspace.testProject(symbolProvider)
5555
project.moduleFor(shape) {
56-
ClientEnumGenerator(codegenContext, shape).render(this)
56+
ClientEnumGenerator(codegenContext, shape, emptyList()).render(this)
5757
unitTest("generate_named_enums") {
5858
withBlock("let result = ", ";") {
5959
sut.render(this, shape, data)
@@ -74,7 +74,7 @@ internal class ClientInstantiatorTest {
7474

7575
val project = TestWorkspace.testProject(symbolProvider)
7676
project.moduleFor(shape) {
77-
ClientEnumGenerator(codegenContext, shape).render(this)
77+
ClientEnumGenerator(codegenContext, shape, emptyList()).render(this)
7878
unitTest("generate_unnamed_enums") {
7979
withBlock("let result = ", ";") {
8080
sut.render(this, shape, data)

‎codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt‎

Lines changed: 11 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.ModuleDocProvider
1212
import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
1313
import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
1414
import software.amazon.smithy.rust.codegen.core.smithy.generators.BuilderCustomization
15+
import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumCustomization
1516
import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization
1617
import software.amazon.smithy.rust.codegen.core.smithy.generators.ManifestCustomizations
1718
import software.amazon.smithy.rust.codegen.core.smithy.generators.StructureCustomization
@@ -59,7 +60,8 @@ interface CoreCodegenDecorator<CodegenContext, CodegenSettings> {
5960
fun extras(
6061
codegenContext: CodegenContext,
6162
rustCrate: RustCrate,
62-
) {}
63+
) {
64+
}
6365

6466
/**
6567
* Customize the documentation provider for module documentation.
@@ -94,6 +96,14 @@ interface CoreCodegenDecorator<CodegenContext, CodegenSettings> {
9496
baseCustomizations: List<StructureCustomization>,
9597
): List<StructureCustomization> = baseCustomizations
9698

99+
/**
100+
* Hook to customize enums generated by `EnumGenerator`.
101+
*/
102+
fun enumCustomizations(
103+
codegenContext: CodegenContext,
104+
baseCustomizations: List<EnumCustomization>,
105+
): List<EnumCustomization> = baseCustomizations
106+
97107
// TODO(https://github.com/smithy-lang/smithy-rs/issues/1401): Move builder customizations into `ClientCodegenDecorator`
98108

99109
/**

‎codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt‎

Lines changed: 36 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,9 @@ import software.amazon.smithy.rust.codegen.core.smithy.MaybeRenamed
2929
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType
3030
import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope
3131
import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
32+
import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization
33+
import software.amazon.smithy.rust.codegen.core.smithy.customize.Section
34+
import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations
3235
import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata
3336
import software.amazon.smithy.rust.codegen.core.smithy.renamedFrom
3437
import software.amazon.smithy.rust.codegen.core.util.REDACTION
@@ -39,6 +42,25 @@ import software.amazon.smithy.rust.codegen.core.util.orNull
3942
import software.amazon.smithy.rust.codegen.core.util.shouldRedact
4043
import software.amazon.smithy.rust.codegen.core.util.toPascalCase
4144

45+
/** EnumGenerator customization sections */
46+
sealed class EnumSection(name: String) : Section(name) {
47+
abstract val shape: Shape
48+
49+
/** Hook to add additional attributes to an enum member */
50+
data class AdditionalMemberAttributes(override val shape: Shape, val definition: EnumDefinition) :
51+
EnumSection("AdditionalMemberAttributes")
52+
53+
/** Hook to add additional trait implementations */
54+
data class AdditionalTraitImpls(override val shape: Shape) : EnumSection("AdditionalTraitImpls")
55+
56+
/** Hook to add additional enum members */
57+
data class AdditionalEnumMembers(override val shape: Shape) :
58+
EnumSection("AdditionalEnumMembers")
59+
}
60+
61+
/** Customizations for EnumGenerator */
62+
abstract class EnumCustomization : NamedCustomization<EnumSection>()
63+
4264
data class EnumGeneratorContext(
4365
val enumName: String,
4466
val enumMeta: RustMetadata,
@@ -86,6 +108,7 @@ class EnumMemberModel(
86108
private val parentShape: Shape,
87109
private val definition: EnumDefinition,
88110
private val symbolProvider: RustSymbolProvider,
111+
private val customizations: List<EnumCustomization>,
89112
) {
90113
companion object {
91114
/**
@@ -140,6 +163,10 @@ class EnumMemberModel(
140163
fun render(writer: RustWriter) {
141164
renderDocumentation(writer)
142165
renderDeprecated(writer)
166+
writer.writeCustomizations(
167+
customizations,
168+
EnumSection.AdditionalMemberAttributes(parentShape, definition),
169+
)
143170
writer.write("${derivedName()},")
144171
}
145172
}
@@ -167,6 +194,7 @@ open class EnumGenerator(
167194
private val symbolProvider: RustSymbolProvider,
168195
private val shape: StringShape,
169196
private val enumType: EnumType,
197+
private val customizations: List<EnumCustomization>,
170198
) {
171199
companion object {
172200
/** Name of the function on the enum impl to get a vec of value names */
@@ -180,7 +208,9 @@ open class EnumGenerator(
180208
enumName = symbol.name,
181209
enumMeta = symbol.expectRustMetadata(),
182210
enumTrait = enumTrait,
183-
sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(shape, it, symbolProvider) },
211+
sortedMembers =
212+
enumTrait.values.sortedBy { it.value }
213+
.map { EnumMemberModel(shape, it, symbolProvider, customizations) },
184214
)
185215

186216
fun render(writer: RustWriter) {
@@ -193,6 +223,7 @@ open class EnumGenerator(
193223
writer.renderUnnamedEnum()
194224
}
195225
enumType.additionalEnumImpls(context)(writer)
226+
writer.writeCustomizations(customizations, EnumSection.AdditionalTraitImpls(shape))
196227

197228
if (shape.shouldRedact(model)) {
198229
writer.renderDebugImplForSensitiveEnum()
@@ -266,6 +297,10 @@ open class EnumGenerator(
266297
context.enumMeta.render(this)
267298
rustBlock("enum ${context.enumName}") {
268299
context.sortedMembers.forEach { member -> member.render(this) }
300+
writeCustomizations(
301+
customizations,
302+
EnumSection.AdditionalEnumMembers(shape),
303+
)
269304
enumType.additionalEnumMembers(context)(this)
270305
}
271306
}

‎codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt‎

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -65,6 +65,7 @@ class EnumGeneratorTest {
6565
testModel.lookup("test#EnumWithUnknown"),
6666
enumTrait.values.first { it.name.orNull() == name },
6767
symbolProvider,
68+
emptyList(),
6869
)
6970

7071
@Test
@@ -112,7 +113,7 @@ class EnumGeneratorTest {
112113
shape: StringShape,
113114
enumType: EnumType = TestEnumType,
114115
) {
115-
EnumGenerator(model, provider, shape, enumType).render(this)
116+
EnumGenerator(model, provider, shape, enumType, emptyList()).render(this)
116117
}
117118

118119
@Test

‎codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/JsonParserGeneratorTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -227,7 +227,7 @@ class JsonParserGeneratorTest {
227227
project.moduleFor(top) {
228228
UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render()
229229
val enum = model.lookup<StringShape>("test#FooEnum")
230-
EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
230+
EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
231231
}
232232
}
233233
model.lookup<OperationShape>("test#Op").outputShape(model).also { output ->

‎codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/XmlBindingTraitParserGeneratorTest.kt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -208,7 +208,7 @@ internal class XmlBindingTraitParserGeneratorTest {
208208
project.moduleFor(top) {
209209
UnionGenerator(model, symbolProvider, this, choiceShape).render()
210210
model.lookup<StringShape>("test#FooEnum").also { enum ->
211-
EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
211+
EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
212212
}
213213
}
214214
}

0 commit comments

Comments
 (0)