Merge branch 'main' into feature/flexible-checksums

smithy-lang · Dec 2, 2024 · 4b25bdd · 4b25bdd
2 parents b181637 + 65035d5
commit 4b25bdd
Show file tree

Hide file tree

Showing 4 changed files with 68 additions and 11 deletions.
diff --git a/...src/main/kotlin/software/amazon/smithy/rust/codegen/client/smithy/ClientCodegenVisitor.kt b/...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/...in/software/amazon/smithy/rust/codegen/client/smithy/generators/OperationCustomization.kt b/...in/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<OperationCustomization>,
+        val operationShape: OperationShape,
+    ) : OperationSection("CustomizableOperationImpl")
 }
 
 abstract class OperationCustomization : NamedCustomization<OperationSection>()
diff --git a/...smithy/rust/codegen/client/smithy/generators/client/CustomizableOperationImplGenerator.kt b/...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<OperationCustomization>,
+) {
+    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<E, B> CustomizableOperation<#T, E, B> {", codegenContext.symbolProvider.toSymbol(operation))
+        writer.writeCustomizations(customizations, section)
+        writer.rust("}")
+    }
+}
diff --git a/...rc/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt b/...rc/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/customize/Customization.kt
@@ -80,3 +80,13 @@ fun <T : Section> RustWriter.writeCustomizationsOrElse(
         orElse(this)
     }
 }
+
+fun <T : Section> allCustomizationsAreEmpty(
+    customizations: List<NamedCustomization<T>>,
+    section: T,
+): Boolean {
+    val test = RustWriter.root()
+    test.writeCustomizations(customizations, section)
+    // If they're not dirty, then they're empty.
+    return !test.dirty()
+}