From 39a81677283f66927d5c6655d20b5cc251be5920 Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Thu, 27 Feb 2025 15:56:37 -0800 Subject: [PATCH 1/7] feat/add-enum-customization-for-member-attributes --- .../core/smithy/generators/EnumGenerator.kt | 19 ++++++++++++++++++- .../smithy/generators/EnumGeneratorTest.kt | 3 ++- .../parse/JsonParserGeneratorTest.kt | 2 +- .../XmlBindingTraitParserGeneratorTest.kt | 2 +- .../AwsQuerySerializerGeneratorTest.kt | 4 ++-- .../Ec2QuerySerializerGeneratorTest.kt | 4 ++-- .../serialize/JsonSerializerGeneratorTest.kt | 4 ++-- .../XmlBindingTraitSerializerGeneratorTest.kt | 4 ++-- 8 files changed, 30 insertions(+), 12 deletions(-) diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt index 02b0ba14d7..cf25f00083 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt @@ -29,6 +29,9 @@ import software.amazon.smithy.rust.codegen.core.smithy.MaybeRenamed import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider +import software.amazon.smithy.rust.codegen.core.smithy.customize.NamedCustomization +import software.amazon.smithy.rust.codegen.core.smithy.customize.Section +import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations import software.amazon.smithy.rust.codegen.core.smithy.expectRustMetadata import software.amazon.smithy.rust.codegen.core.smithy.renamedFrom import software.amazon.smithy.rust.codegen.core.util.REDACTION @@ -39,6 +42,17 @@ import software.amazon.smithy.rust.codegen.core.util.orNull import software.amazon.smithy.rust.codegen.core.util.shouldRedact import software.amazon.smithy.rust.codegen.core.util.toPascalCase +/** EnumGenerator customization sections */ +sealed class EnumSection(name: String) : Section(name) { + abstract val definition: EnumDefinition + + /** Hook to add additional attributes to an enum member */ + data class AdditionalMemberAttributes(override val definition: EnumDefinition) : EnumSection("AdditionalMemberAttributes") +} + +/** Customizations for EnumGenerator */ +abstract class EnumCustomization : NamedCustomization() + data class EnumGeneratorContext( val enumName: String, val enumMeta: RustMetadata, @@ -86,6 +100,7 @@ class EnumMemberModel( private val parentShape: Shape, private val definition: EnumDefinition, private val symbolProvider: RustSymbolProvider, + private val customizations: List, ) { companion object { /** @@ -140,6 +155,7 @@ class EnumMemberModel( fun render(writer: RustWriter) { renderDocumentation(writer) renderDeprecated(writer) + writer.writeCustomizations(customizations, EnumSection.AdditionalMemberAttributes(definition)) writer.write("${derivedName()},") } } @@ -167,6 +183,7 @@ open class EnumGenerator( private val symbolProvider: RustSymbolProvider, private val shape: StringShape, private val enumType: EnumType, + private val customizations: List, ) { companion object { /** Name of the function on the enum impl to get a vec of value names */ @@ -180,7 +197,7 @@ open class EnumGenerator( enumName = symbol.name, enumMeta = symbol.expectRustMetadata(), enumTrait = enumTrait, - sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(shape, it, symbolProvider) }, + sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(shape, it, symbolProvider, customizations) }, ) fun render(writer: RustWriter) { diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt index 0528c2e364..d5fe0e7e3f 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGeneratorTest.kt @@ -65,6 +65,7 @@ class EnumGeneratorTest { testModel.lookup("test#EnumWithUnknown"), enumTrait.values.first { it.name.orNull() == name }, symbolProvider, + emptyList(), ) @Test @@ -112,7 +113,7 @@ class EnumGeneratorTest { shape: StringShape, enumType: EnumType = TestEnumType, ) { - EnumGenerator(model, provider, shape, enumType).render(this) + EnumGenerator(model, provider, shape, enumType, emptyList()).render(this) } @Test diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/JsonParserGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/JsonParserGeneratorTest.kt index 8e7985a27a..0a5c8376c4 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/JsonParserGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/JsonParserGeneratorTest.kt @@ -227,7 +227,7 @@ class JsonParserGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } model.lookup("test#Op").outputShape(model).also { output -> diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/XmlBindingTraitParserGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/XmlBindingTraitParserGeneratorTest.kt index 0ec685f15a..9099ebb9cb 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/XmlBindingTraitParserGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/parse/XmlBindingTraitParserGeneratorTest.kt @@ -208,7 +208,7 @@ internal class XmlBindingTraitParserGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, choiceShape).render() model.lookup("test#FooEnum").also { enum -> - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } } diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/AwsQuerySerializerGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/AwsQuerySerializerGeneratorTest.kt index a0fd7fffe8..aaf483fb52 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/AwsQuerySerializerGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/AwsQuerySerializerGeneratorTest.kt @@ -147,7 +147,7 @@ class AwsQuerySerializerGeneratorTest { renderUnknownVariant = generateUnknownVariant, ).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } @@ -316,7 +316,7 @@ class AwsQuerySerializerGeneratorTest { renderUnknownVariant = generateUnknownVariant, ).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/Ec2QuerySerializerGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/Ec2QuerySerializerGeneratorTest.kt index d0a33e1f5d..5012daba66 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/Ec2QuerySerializerGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/Ec2QuerySerializerGeneratorTest.kt @@ -143,7 +143,7 @@ class Ec2QuerySerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } @@ -298,7 +298,7 @@ class Ec2QuerySerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/JsonSerializerGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/JsonSerializerGeneratorTest.kt index 922f2b7ae8..471a640d08 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/JsonSerializerGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/JsonSerializerGeneratorTest.kt @@ -159,7 +159,7 @@ class JsonSerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } @@ -333,7 +333,7 @@ class JsonSerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } diff --git a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/XmlBindingTraitSerializerGeneratorTest.kt b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/XmlBindingTraitSerializerGeneratorTest.kt index 58f482e276..e612b72da2 100644 --- a/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/XmlBindingTraitSerializerGeneratorTest.kt +++ b/codegen-core/src/test/kotlin/software/amazon/smithy/rust/codegen/core/smithy/protocols/serialize/XmlBindingTraitSerializerGeneratorTest.kt @@ -164,7 +164,7 @@ internal class XmlBindingTraitSerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } model.lookup("test#Op").inputShape(model).also { input -> @@ -334,7 +334,7 @@ internal class XmlBindingTraitSerializerGeneratorTest { project.moduleFor(top) { UnionGenerator(model, symbolProvider, this, model.lookup("test#Choice")).render() val enum = model.lookup("test#FooEnum") - EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this) + EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this) } } model.lookup("test#Op").inputShape(model).also { input -> From e4bf00ccf16717a4c89b8b65dcc8dda5fb6602da Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Thu, 27 Feb 2025 16:59:41 -0800 Subject: [PATCH 2/7] chore: add changelog --- .changelog/1740703869.md | 8 ++++++++ 1 file changed, 8 insertions(+) create mode 100644 .changelog/1740703869.md diff --git a/.changelog/1740703869.md b/.changelog/1740703869.md new file mode 100644 index 0000000000..c4aef5a55a --- /dev/null +++ b/.changelog/1740703869.md @@ -0,0 +1,8 @@ +--- +applies_to: [both] +authors: [Dorenavant] +references: [] +breaking: false +new_feature: true +bug_fix: false +--- From 1df9a0e19bd119ad74d6dc1ba4f1ee9add2f2664 Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Thu, 27 Mar 2025 14:59:28 -0700 Subject: [PATCH 3/7] Add AdditionalTraitImpls hook --- .../core/smithy/generators/EnumGenerator.kt | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt index cf25f00083..a1c384d779 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt @@ -7,6 +7,7 @@ package software.amazon.smithy.rust.codegen.core.smithy.generators import software.amazon.smithy.codegen.core.Symbol import software.amazon.smithy.model.Model +import software.amazon.smithy.model.shapes.EnumShape import software.amazon.smithy.model.shapes.MemberShape import software.amazon.smithy.model.shapes.Shape import software.amazon.smithy.model.shapes.StringShape @@ -44,10 +45,13 @@ import software.amazon.smithy.rust.codegen.core.util.toPascalCase /** EnumGenerator customization sections */ sealed class EnumSection(name: String) : Section(name) { - abstract val definition: EnumDefinition + abstract val shape: Shape /** Hook to add additional attributes to an enum member */ - data class AdditionalMemberAttributes(override val definition: EnumDefinition) : EnumSection("AdditionalMemberAttributes") + data class AdditionalMemberAttributes(override val shape: Shape, val definition: EnumDefinition) : + EnumSection("AdditionalMemberAttributes") + + data class AdditionalTraitImpls(override val shape: Shape) : EnumSection("AdditionalTraitImpls") } /** Customizations for EnumGenerator */ @@ -155,7 +159,10 @@ class EnumMemberModel( fun render(writer: RustWriter) { renderDocumentation(writer) renderDeprecated(writer) - writer.writeCustomizations(customizations, EnumSection.AdditionalMemberAttributes(definition)) + writer.writeCustomizations( + customizations, + EnumSection.AdditionalMemberAttributes(parentShape, definition), + ) writer.write("${derivedName()},") } } @@ -197,7 +204,8 @@ open class EnumGenerator( enumName = symbol.name, enumMeta = symbol.expectRustMetadata(), enumTrait = enumTrait, - sortedMembers = enumTrait.values.sortedBy { it.value }.map { EnumMemberModel(shape, it, symbolProvider, customizations) }, + sortedMembers = enumTrait.values.sortedBy { it.value } + .map { EnumMemberModel(shape, it, symbolProvider, customizations) }, ) fun render(writer: RustWriter) { @@ -210,6 +218,7 @@ open class EnumGenerator( writer.renderUnnamedEnum() } enumType.additionalEnumImpls(context)(writer) + writer.writeCustomizations(customizations, EnumSection.AdditionalTraitImpls(shape)) if (shape.shouldRedact(model)) { writer.renderDebugImplForSensitiveEnum() From 3a0045d63f970dfa92c5b08a576b195d087d0280 Mon Sep 17 00:00:00 2001 From: Landon James Date: Mon, 31 Mar 2025 13:00:36 -0700 Subject: [PATCH 4/7] Update changelog so it will parse correctly --- .changelog/1740703869.md | 4 +++- 1 file changed, 3 insertions(+), 1 deletion(-) diff --git a/.changelog/1740703869.md b/.changelog/1740703869.md index c4aef5a55a..0591f74eea 100644 --- a/.changelog/1740703869.md +++ b/.changelog/1740703869.md @@ -1,8 +1,10 @@ --- -applies_to: [both] +applies_to: ["client", "server"] authors: [Dorenavant] references: [] breaking: false new_feature: true bug_fix: false --- + +Add EnumSection to allow decorators to modify enum member attributes From f71c8c4531ab45a322cce2691c644fc0976e4167 Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Tue, 1 Apr 2025 12:17:59 -0700 Subject: [PATCH 5/7] fix: integrate decorator with enumCustomizations --- .../client/smithy/ClientCodegenVisitor.kt | 25 ++++++++++++++++--- .../smithy/generators/ClientEnumGenerator.kt | 7 +++++- .../smithy/customize/CoreCodegenDecorator.kt | 12 ++++++++- .../server/smithy/ServerCodegenVisitor.kt | 7 +++++- .../smithy/generators/ServerEnumGenerator.kt | 16 +++++++----- 5 files changed, 54 insertions(+), 13 deletions(-) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt index ae1788d946..1403049023 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt @@ -93,7 +93,13 @@ class ClientCodegenVisitor( model = codegenDecorator.transformModel(untransformedService, baseModel, settings) // the model transformer _might_ change the service shape val service = settings.getService(model) - symbolProvider = RustClientCodegenPlugin.baseSymbolProvider(settings, model, service, rustSymbolProviderConfig, codegenDecorator) + symbolProvider = RustClientCodegenPlugin.baseSymbolProvider( + settings, + model, + service, + rustSymbolProviderConfig, + codegenDecorator, + ) codegenContext = ClientCodegenContext( @@ -177,7 +183,10 @@ class ClientCodegenVisitor( ) try { // use an increased max_width to make rustfmt fail less frequently - "cargo fmt -- --config max_width=150".runCommand(fileManifest.baseDir, timeout = settings.codegenConfig.formatTimeoutSeconds.toLong()) + "cargo fmt -- --config max_width=150".runCommand( + fileManifest.baseDir, + timeout = settings.codegenConfig.formatTimeoutSeconds.toLong(), + ) } catch (err: CommandError) { logger.warning("Failed to run cargo fmt: [${service.id}]\n${err.output}") } @@ -236,7 +245,10 @@ class ClientCodegenVisitor( implBlock(symbolProvider.toSymbol(shape)) { BuilderGenerator.renderConvenienceMethod(this, symbolProvider, shape) - if (codegenContext.protocolImpl?.httpBindingResolver?.handlesEventStreamInitialResponse(shape) == true) { + if (codegenContext.protocolImpl?.httpBindingResolver?.handlesEventStreamInitialResponse( + shape, + ) == true + ) { BuilderGenerator.renderIntoBuilderMethod(this, symbolProvider, shape) } } @@ -251,6 +263,7 @@ class ClientCodegenVisitor( } struct to builder } + else -> { val errorGenerator = ErrorGenerator( @@ -283,7 +296,11 @@ class ClientCodegenVisitor( if (shape.hasTrait()) { val privateModule = privateModule(shape) rustCrate.inPrivateModuleWithReexport(privateModule, symbolProvider.toSymbol(shape)) { - ClientEnumGenerator(codegenContext, shape).render(this) + ClientEnumGenerator( + codegenContext, + shape, + codegenDecorator.enumCustomizations(codegenContext, emptyList()), + ).render(this) } } } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt index f5bbbefb9c..5726b267b2 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt @@ -19,6 +19,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate import software.amazon.smithy.rust.codegen.core.rustlang.writable import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope +import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGenerator import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGeneratorContext import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumMemberModel @@ -278,7 +279,10 @@ data class InfallibleEnumType( } } -class ClientEnumGenerator(codegenContext: ClientCodegenContext, shape: StringShape) : +class ClientEnumGenerator( + codegenContext: ClientCodegenContext, shape: StringShape, + customizations: List, +) : EnumGenerator( codegenContext.model, codegenContext.symbolProvider, @@ -290,6 +294,7 @@ class ClientEnumGenerator(codegenContext: ClientCodegenContext, shape: StringSha parent = ClientRustModule.primitives, ), ), + customizations, ) private fun unknownVariantError(): RuntimeType = diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt index fa83c62208..ee66980214 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/CoreCodegenDecorator.kt @@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.ModuleDocProvider import software.amazon.smithy.rust.codegen.core.smithy.RustCrate import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider import software.amazon.smithy.rust.codegen.core.smithy.generators.BuilderCustomization +import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.LibRsCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.ManifestCustomizations import software.amazon.smithy.rust.codegen.core.smithy.generators.StructureCustomization @@ -59,7 +60,8 @@ interface CoreCodegenDecorator { fun extras( codegenContext: CodegenContext, rustCrate: RustCrate, - ) {} + ) { + } /** * Customize the documentation provider for module documentation. @@ -94,6 +96,14 @@ interface CoreCodegenDecorator { baseCustomizations: List, ): List = baseCustomizations + /** + * Hook to customize enums generated by `EnumGenerator`. + */ + fun enumCustomizations( + codegenContext: CodegenContext, + baseCustomizations: List, + ): List = baseCustomizations + // TODO(https://github.com/smithy-lang/smithy-rs/issues/1401): Move builder customizations into `ClientCodegenDecorator` /** diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt index b1fc07a31a..a4a8b5afff 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/ServerCodegenVisitor.kt @@ -512,7 +512,12 @@ open class ServerCodegenVisitor( fun serverEnumGeneratorFactory( codegenContext: ServerCodegenContext, shape: StringShape, - ) = ServerEnumGenerator(codegenContext, shape, validationExceptionConversionGenerator) + ) = ServerEnumGenerator( + codegenContext, + shape, + validationExceptionConversionGenerator, + codegenDecorator.enumCustomizations(codegenContext, emptyList()), + ) stringShape(shape, ::serverEnumGeneratorFactory) } diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt index a3cc269692..76946f4d61 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt @@ -12,6 +12,7 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate import software.amazon.smithy.rust.codegen.core.rustlang.writable import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType import software.amazon.smithy.rust.codegen.core.smithy.RuntimeType.Companion.preludeScope +import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumCustomization import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGenerator import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumGeneratorContext import software.amazon.smithy.rust.codegen.core.smithy.generators.EnumType @@ -51,7 +52,9 @@ open class ConstrainedEnum( impl #{Display} for $constraintViolationName { fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { - write!(f, r##"${context.enumTrait.shapeConstraintViolationDisplayMessage(shape).replace("#", "##")}"##) + write!(f, r##"${ + context.enumTrait.shapeConstraintViolationDisplayMessage(shape).replace("#", "##") + }"##) } } @@ -162,9 +165,10 @@ class ServerEnumGenerator( codegenContext: ServerCodegenContext, shape: StringShape, validationExceptionConversionGenerator: ValidationExceptionConversionGenerator, + customizations: List, ) : EnumGenerator( - codegenContext.model, - codegenContext.symbolProvider, - shape, - enumType = ConstrainedEnum(codegenContext, shape, validationExceptionConversionGenerator), - ) + codegenContext.model, + codegenContext.symbolProvider, + shape, + enumType = ConstrainedEnum(codegenContext, shape, validationExceptionConversionGenerator), customizations, +) From f6148026765124a957aee6d49b14634f58509e72 Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Mon, 14 Apr 2025 07:52:56 -0700 Subject: [PATCH 6/7] feat: add additional enum members hook --- .../codegen/core/smithy/generators/EnumGenerator.kt | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt index a1c384d779..d56defdcff 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/generators/EnumGenerator.kt @@ -7,7 +7,6 @@ package software.amazon.smithy.rust.codegen.core.smithy.generators import software.amazon.smithy.codegen.core.Symbol import software.amazon.smithy.model.Model -import software.amazon.smithy.model.shapes.EnumShape import software.amazon.smithy.model.shapes.MemberShape import software.amazon.smithy.model.shapes.Shape import software.amazon.smithy.model.shapes.StringShape @@ -51,7 +50,12 @@ sealed class EnumSection(name: String) : Section(name) { data class AdditionalMemberAttributes(override val shape: Shape, val definition: EnumDefinition) : EnumSection("AdditionalMemberAttributes") + /** Hook to add additional trait implementations */ data class AdditionalTraitImpls(override val shape: Shape) : EnumSection("AdditionalTraitImpls") + + /** Hook to add additional enum members */ + data class AdditionalEnumMembers(override val shape: Shape) : + EnumSection("AdditionalEnumMembers") } /** Customizations for EnumGenerator */ @@ -292,6 +296,10 @@ open class EnumGenerator( context.enumMeta.render(this) rustBlock("enum ${context.enumName}") { context.sortedMembers.forEach { member -> member.render(this) } + writeCustomizations( + customizations, + EnumSection.AdditionalEnumMembers(shape), + ) enumType.additionalEnumMembers(context)(this) } } From 6ad5741bf45ef5a8b3ec692b99e33ad18efb87ff Mon Sep 17 00:00:00 2001 From: Ignatius Sinn Date: Mon, 14 Apr 2025 07:54:33 -0700 Subject: [PATCH 7/7] chore: fix tests --- .../client/smithy/generators/ClientEnumGenerator.kt | 3 ++- .../client/smithy/generators/ClientEnumGeneratorTest.kt | 8 ++++---- .../client/smithy/generators/ClientInstantiatorTest.kt | 4 ++-- .../server/smithy/generators/ServerEnumGenerator.kt | 3 ++- .../smithy/generators/ServerBuilderDefaultValuesTest.kt | 3 +++ .../server/smithy/generators/ServerEnumGeneratorTest.kt | 3 +++ .../server/smithy/generators/ServerInstantiatorTest.kt | 2 ++ 7 files changed, 18 insertions(+), 8 deletions(-) diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt index 5726b267b2..d230a0f78f 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGenerator.kt @@ -280,7 +280,8 @@ data class InfallibleEnumType( } class ClientEnumGenerator( - codegenContext: ClientCodegenContext, shape: StringShape, + codegenContext: ClientCodegenContext, + shape: StringShape, customizations: List, ) : EnumGenerator( diff --git a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGeneratorTest.kt b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGeneratorTest.kt index 15ac5c41fc..457c85dffd 100644 --- a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGeneratorTest.kt +++ b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientEnumGeneratorTest.kt @@ -28,7 +28,7 @@ class ClientEnumGeneratorTest { val context = testClientCodegenContext(model) val project = TestWorkspace.testProject(context.symbolProvider) project.moduleFor(shape) { - ClientEnumGenerator(context, shape).render(this) + ClientEnumGenerator(context, shape, emptyList()).render(this) unitTest( "matching_on_enum_should_be_forward_compatible", """ @@ -88,7 +88,7 @@ class ClientEnumGeneratorTest { val context = testClientCodegenContext(model) val project = TestWorkspace.testProject(context.symbolProvider) project.moduleFor(shape) { - ClientEnumGenerator(context, shape).render(this) + ClientEnumGenerator(context, shape, emptyList()).render(this) unitTest( "impl_debug_for_non_sensitive_enum_should_implement_the_derived_debug_trait", """ @@ -123,7 +123,7 @@ class ClientEnumGeneratorTest { val context = testClientCodegenContext(model) val project = TestWorkspace.testProject(context.symbolProvider) project.moduleFor(shape) { - ClientEnumGenerator(context, shape).render(this) + ClientEnumGenerator(context, shape, emptyList()).render(this) unitTest( "it_escapes_the_unknown_variant_if_the_enum_has_an_unknown_value_in_the_model", """ @@ -156,7 +156,7 @@ class ClientEnumGeneratorTest { val project = TestWorkspace.testProject(context.symbolProvider) project.moduleFor(shape) { rust("##![allow(deprecated)]") - ClientEnumGenerator(context, shape).render(this) + ClientEnumGenerator(context, shape, emptyList()).render(this) unitTest( "generated_named_enums_roundtrip", """ diff --git a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientInstantiatorTest.kt b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientInstantiatorTest.kt index 60ee794914..990e3a90d2 100644 --- a/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientInstantiatorTest.kt +++ b/codegen-client/src/test/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/ClientInstantiatorTest.kt @@ -53,7 +53,7 @@ internal class ClientInstantiatorTest { val project = TestWorkspace.testProject(symbolProvider) project.moduleFor(shape) { - ClientEnumGenerator(codegenContext, shape).render(this) + ClientEnumGenerator(codegenContext, shape, emptyList()).render(this) unitTest("generate_named_enums") { withBlock("let result = ", ";") { sut.render(this, shape, data) @@ -74,7 +74,7 @@ internal class ClientInstantiatorTest { val project = TestWorkspace.testProject(symbolProvider) project.moduleFor(shape) { - ClientEnumGenerator(codegenContext, shape).render(this) + ClientEnumGenerator(codegenContext, shape, emptyList()).render(this) unitTest("generate_unnamed_enums") { withBlock("let result = ", ";") { sut.render(this, shape, data) diff --git a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt index 76946f4d61..68595e6a77 100644 --- a/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt +++ b/codegen-server/src/main/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGenerator.kt @@ -170,5 +170,6 @@ class ServerEnumGenerator( codegenContext.model, codegenContext.symbolProvider, shape, - enumType = ConstrainedEnum(codegenContext, shape, validationExceptionConversionGenerator), customizations, + enumType = ConstrainedEnum(codegenContext, shape, validationExceptionConversionGenerator), + customizations, ) diff --git a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerBuilderDefaultValuesTest.kt b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerBuilderDefaultValuesTest.kt index e0a2054285..0d6cdd0476 100644 --- a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerBuilderDefaultValuesTest.kt +++ b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerBuilderDefaultValuesTest.kt @@ -112,6 +112,7 @@ class ServerBuilderDefaultValuesTest { BuilderGeneratorKind.SERVER_BUILDER_GENERATOR -> { writeServerBuilderGenerator(project, this, model, symbolProvider) } + BuilderGeneratorKind.SERVER_BUILDER_GENERATOR_WITHOUT_PUBLIC_CONSTRAINED_TYPES -> { writeServerBuilderGeneratorWithoutPublicConstrainedTypes(project, this, model, symbolProvider) } @@ -211,6 +212,7 @@ class ServerBuilderDefaultValuesTest { codegenContext, model.lookup("com.test#Language"), SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(writer) StructureGenerator(model, symbolProvider, writer, struct, emptyList(), codegenContext.structSettings()).render() } @@ -240,6 +242,7 @@ class ServerBuilderDefaultValuesTest { codegenContext, model.lookup("com.test#Language"), SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(writer) StructureGenerator(model, symbolProvider, writer, struct, emptyList(), codegenContext.structSettings()).render() } diff --git a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGeneratorTest.kt b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGeneratorTest.kt index 3236106032..fb05a8f011 100644 --- a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGeneratorTest.kt +++ b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerEnumGeneratorTest.kt @@ -46,6 +46,7 @@ class ServerEnumGeneratorTest { codegenContext, shape, SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(writer) writer.compileAndTest( """ @@ -63,6 +64,7 @@ class ServerEnumGeneratorTest { codegenContext, shape, SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(writer) writer.compileAndTest( """ @@ -82,6 +84,7 @@ class ServerEnumGeneratorTest { codegenContext, shape, SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(writer) writer.toString() shouldNotContain "#[non_exhaustive]" } diff --git a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerInstantiatorTest.kt b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerInstantiatorTest.kt index 1a1f9f01ea..8c816cf0a3 100644 --- a/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerInstantiatorTest.kt +++ b/codegen-server/src/test/kotlin/software/amazon/smithy/rust/codegen/server/smithy/generators/ServerInstantiatorTest.kt @@ -214,6 +214,7 @@ class ServerInstantiatorTest { codegenContext, shape, SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(this) unitTest("generate_named_enums") { withBlock("let result = ", ";") { @@ -237,6 +238,7 @@ class ServerInstantiatorTest { codegenContext, shape, SmithyValidationExceptionConversionGenerator(codegenContext), + emptyList(), ).render(this) unitTest("generate_unnamed_enums") { withBlock("let result = ", ";") {