Add EnumSection to allow decorators to modify enum member attributes

.changelog/1740703869.md

@@ -0,0 +1,10 @@
+---
+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

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<EnumTrait>()) {
             val privateModule = privateModule(shape)
             rustCrate.inPrivateModuleWithReexport(privateModule, symbolProvider.toSymbol(shape)) {
-                ClientEnumGenerator(codegenContext, shape).render(this)
+                ClientEnumGenerator(
+                    codegenContext,
+                    shape,
+                    codegenDecorator.enumCustomizations(codegenContext, emptyList()),
+                ).render(this)
             }
         }
     }

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,11 @@ data class InfallibleEnumType(
     }
 }
 
-class ClientEnumGenerator(codegenContext: ClientCodegenContext, shape: StringShape) :
+class ClientEnumGenerator(
+    codegenContext: ClientCodegenContext,
+    shape: StringShape,
+    customizations: List<EnumCustomization>,
+) :
     EnumGenerator(
         codegenContext.model,
         codegenContext.symbolProvider,
@@ -290,6 +295,7 @@ class ClientEnumGenerator(codegenContext: ClientCodegenContext, shape: StringSha
                 parent = ClientRustModule.primitives,
             ),
         ),
+        customizations,
     )
 
 private fun unknownVariantError(): RuntimeType =

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",
                 """

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)

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<CodegenContext, CodegenSettings> {
     fun extras(
         codegenContext: CodegenContext,
         rustCrate: RustCrate,
-    ) {}
+    ) {
+    }
 
     /**
      * Customize the documentation provider for module documentation.
@@ -94,6 +96,14 @@ interface CoreCodegenDecorator<CodegenContext, CodegenSettings> {
         baseCustomizations: List<StructureCustomization>,
     ): List<StructureCustomization> = baseCustomizations
 
+    /**
+     * Hook to customize enums generated by `EnumGenerator`.
+     */
+    fun enumCustomizations(
+        codegenContext: CodegenContext,
+        baseCustomizations: List<EnumCustomization>,
+    ): List<EnumCustomization> = baseCustomizations
+
     // TODO(https://github.com/smithy-lang/smithy-rs/issues/1401): Move builder customizations into `ClientCodegenDecorator`
 
     /**

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,25 @@ 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 shape: Shape
+
+    /** Hook to add additional attributes to an enum member */
+    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 */
+abstract class EnumCustomization : NamedCustomization<EnumSection>()
+
 data class EnumGeneratorContext(
     val enumName: String,
     val enumMeta: RustMetadata,
@@ -86,6 +108,7 @@ class EnumMemberModel(
     private val parentShape: Shape,
     private val definition: EnumDefinition,
     private val symbolProvider: RustSymbolProvider,
+    private val customizations: List<EnumCustomization>,
 ) {
     companion object {
         /**
@@ -140,6 +163,10 @@ class EnumMemberModel(
     fun render(writer: RustWriter) {
         renderDocumentation(writer)
         renderDeprecated(writer)
+        writer.writeCustomizations(
+            customizations,
+            EnumSection.AdditionalMemberAttributes(parentShape, definition),
+        )
         writer.write("${derivedName()},")
     }
 }
@@ -167,6 +194,7 @@ open class EnumGenerator(
     private val symbolProvider: RustSymbolProvider,
     private val shape: StringShape,
     private val enumType: EnumType,
+    private val customizations: List<EnumCustomization>,
 ) {
     companion object {
         /** Name of the function on the enum impl to get a vec of value names */
@@ -180,7 +208,8 @@ 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) {
@@ -193,6 +222,7 @@ open class EnumGenerator(
             writer.renderUnnamedEnum()
         }
         enumType.additionalEnumImpls(context)(writer)
+        writer.writeCustomizations(customizations, EnumSection.AdditionalTraitImpls(shape))
 
         if (shape.shouldRedact(model)) {
             writer.renderDebugImplForSensitiveEnum()
@@ -266,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)
         }
     }

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

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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
         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

@@ -208,7 +208,7 @@ internal class XmlBindingTraitParserGeneratorTest {
             project.moduleFor(top) {
                 UnionGenerator(model, symbolProvider, this, choiceShape).render()
                 model.lookup<StringShape>("test#FooEnum").also { enum ->
-                    EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                    EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
                 }
             }
         }

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<StringShape>("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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
 

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<StringShape>("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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
 

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<StringShape>("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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
 

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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
         model.lookup<OperationShape>("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<StringShape>("test#FooEnum")
-                EnumGenerator(model, symbolProvider, enum, TestEnumType).render(this)
+                EnumGenerator(model, symbolProvider, enum, TestEnumType, emptyList()).render(this)
             }
         }
         model.lookup<OperationShape>("test#Op").inputShape(model).also { input ->

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)
     }