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 7e8737f50d..24f4e6e38e 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 @@ -21,6 +21,7 @@ import software.amazon.smithy.rust.codegen.client.smithy.customize.ClientCodegen import software.amazon.smithy.rust.codegen.client.smithy.generators.ClientEnumGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.ServiceGenerator +import software.amazon.smithy.rust.codegen.client.smithy.generators.client.CustomizableOperationImplGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.error.ErrorGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.error.OperationErrorGenerator import software.amazon.smithy.rust.codegen.client.smithy.generators.protocol.ClientProtocolTestGenerator @@ -311,10 +312,10 @@ class ClientCodegenVisitor( * Generate operations */ override fun operationShape(operationShape: OperationShape) { - rustCrate.useShapeWriter(operationShape) operationWriter@{ + rustCrate.useShapeWriter(operationShape) { // Render the operation shape operationGenerator.renderOperation( - this@operationWriter, + this, operationShape, codegenDecorator, ) @@ -327,16 +328,24 @@ class ClientCodegenVisitor( protocolGeneratorFactory.support(), operationShape, ), - ).render(this@operationWriter) + ).render(this) + } - rustCrate.withModule(symbolProvider.moduleForOperationError(operationShape)) { - OperationErrorGenerator( - model, - symbolProvider, - operationShape, - codegenDecorator.errorCustomizations(codegenContext, emptyList()), - ).render(this) - } + rustCrate.withModule(symbolProvider.moduleForOperationError(operationShape)) { + OperationErrorGenerator( + model, + symbolProvider, + operationShape, + codegenDecorator.errorCustomizations(codegenContext, emptyList()), + ).render(this) + } + + rustCrate.withModule(ClientRustModule.Client.customize) { + CustomizableOperationImplGenerator( + codegenContext, + operationShape, + codegenDecorator.operationCustomizations(codegenContext, operationShape, emptyList()), + ).render(this) } } } diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt index 07ed4e51cb..2e31dbac44 100644 --- a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt @@ -125,6 +125,11 @@ sealed class OperationSection(name: String) : Section(name) { writer.rustTemplate(".with_retry_classifier(#{classifier})", "classifier" to classifier) } } + + data class CustomizableOperationImpl( + override val customizations: List, + val operationShape: OperationShape, + ) : OperationSection("CustomizableOperationImpl") } abstract class OperationCustomization : NamedCustomization() diff --git a/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt new file mode 100644 index 0000000000..7f839be24b --- /dev/null +++ b/codegen-client/src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt @@ -0,0 +1,33 @@ +/* + * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved. + * SPDX-License-Identifier: Apache-2.0 + */ + +package software.amazon.smithy.rust.codegen.client.smithy.generators.client + +import software.amazon.smithy.model.shapes.OperationShape +import software.amazon.smithy.rust.codegen.client.smithy.ClientCodegenContext +import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationCustomization +import software.amazon.smithy.rust.codegen.client.smithy.generators.OperationSection +import software.amazon.smithy.rust.codegen.core.rustlang.RustWriter +import software.amazon.smithy.rust.codegen.core.rustlang.rust +import software.amazon.smithy.rust.codegen.core.smithy.customize.allCustomizationsAreEmpty +import software.amazon.smithy.rust.codegen.core.smithy.customize.writeCustomizations + +class CustomizableOperationImplGenerator( + private val codegenContext: ClientCodegenContext, + private val operation: OperationShape, + private val customizations: List, +) { + fun render(writer: RustWriter) { + val section = OperationSection.CustomizableOperationImpl(customizations, operation) + // When no customizations are set or there is nothing to write, return early. + if (customizations.isEmpty() || allCustomizationsAreEmpty(customizations, section)) { + return + } + + writer.rust("impl CustomizableOperation<#T, E, B> {", codegenContext.symbolProvider.toSymbol(operation)) + writer.writeCustomizations(customizations, section) + writer.rust("}") + } +} diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt index 79ff1bdb6f..3a969bbb08 100644 --- a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt +++ b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt @@ -80,3 +80,13 @@ fun RustWriter.writeCustomizationsOrElse( orElse(this) } } + +fun allCustomizationsAreEmpty( + customizations: List>, + section: T, +): Boolean { + val test = RustWriter.root() + test.writeCustomizations(customizations, section) + // If they're not dirty, then they're empty. + return !test.dirty() +}