Skip to content

Add EnumSection to allow decorators to modify enum member attributes #4039

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
wants to merge 11 commits into
base: main
Choose a base branch
from
8 changes: 8 additions & 0 deletions .changelog/1740703869.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
---
applies_to: [both]
authors: [Dorenavant]
references: []
breaking: false
new_feature: true
bug_fix: false
---
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand All @@ -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<EnumSection>()

data class EnumGeneratorContext(
val enumName: String,
val enumMeta: RustMetadata,
Expand Down Expand Up @@ -86,6 +100,7 @@ class EnumMemberModel(
private val parentShape: Shape,
private val definition: EnumDefinition,
private val symbolProvider: RustSymbolProvider,
private val customizations: List<EnumCustomization>,
) {
companion object {
/**
Expand Down Expand Up @@ -140,6 +155,7 @@ class EnumMemberModel(
fun render(writer: RustWriter) {
renderDocumentation(writer)
renderDeprecated(writer)
writer.writeCustomizations(customizations, EnumSection.AdditionalMemberAttributes(definition))
writer.write("${derivedName()},")
}
}
Expand Down Expand Up @@ -167,6 +183,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 */
Expand All @@ -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) {
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -65,6 +65,7 @@ class EnumGeneratorTest {
testModel.lookup("test#EnumWithUnknown"),
enumTrait.values.first { it.name.orNull() == name },
symbolProvider,
emptyList(),
)

@Test
Expand Down Expand Up @@ -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
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 ->
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}

Expand Down Expand Up @@ -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)
}
}

Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 ->
Expand Down Expand Up @@ -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 ->
Expand Down