feat: generate operation samples (#1007)

aajtodd · web-flow · commit 5457cfc4d47c · 2023-12-20T10:36:23.000-05:00
diff --git a/.changes/295410ec-cc32-4a9c-a30a-aaeda9996c0f.json b/.changes/295410ec-cc32-4a9c-a30a-aaeda9996c0f.json
@@ -0,0 +1,5 @@
+{
+    "id": "295410ec-cc32-4a9c-a30a-aaeda9996c0f",
+    "type": "feature",
+    "description": "Generate KDoc samples from modeled examples"
+}
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ShapeValueGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ShapeValueGenerator.kt
@@ -25,24 +25,44 @@ class ShapeValueGenerator(
 ) {
 
     /**
-     * Writes generation of a shape value type declaration for the given the parameters.
+     * Renders a shape value declaration for the given parameters inline
+     * with the current writer.
      *
      * @param writer writer to write generated code with.
      * @param shape the shape that will be declared.
      * @param params parameters to fill the generated shape declaration.
      */
-    fun writeShapeValueInline(writer: KotlinWriter, shape: Shape, params: Node) {
-        val nodeVisitor = ShapeValueNodeVisitor(writer, this, shape)
-        when (shape.type) {
-            ShapeType.STRUCTURE -> {
-                if (params.isNullNode) {
-                    params.accept(nodeVisitor)
-                } else {
-                    classDeclaration(writer, shape.asStructureShape().get()) {
-                        params.accept(nodeVisitor)
-                    }
+    fun instantiateShapeInline(writer: KotlinWriter, shape: Shape, params: Node) {
+        if (shape.isStructureShape) {
+            if (params.isNullNode) {
+                writeShapeValuesInline(writer, shape, params)
+            } else {
+                classDeclaration(writer, shape.asStructureShape().get()) {
+                    writeShapeValuesInline(writer, shape, params)
                 }
             }
+        } else {
+            writeShapeValuesInline(writer, shape, params)
+        }
+    }
+
+    /**
+     * Renders the mapping of the shape fields to the given parameters
+     *
+     * @param writer writer to write generated code with.
+     * @param shape the shape that will be declared.
+     * @param params parameters to fill the generated shape declaration.
+     */
+    fun writeShapeValues(writer: KotlinWriter, shape: Shape, params: Node) {
+        writer.ensureNewline()
+        writeShapeValuesInline(writer, shape, params)
+        writer.ensureNewline()
+    }
+
+    private fun writeShapeValuesInline(writer: KotlinWriter, shape: Shape, params: Node) {
+        val nodeVisitor = ShapeValueNodeVisitor(writer, this, shape)
+        when (shape.type) {
+            ShapeType.STRUCTURE -> params.accept(nodeVisitor)
             ShapeType.MAP -> mapDeclaration(writer, shape.asMapShape().get()) {
                 params.accept(nodeVisitor)
             }
@@ -58,11 +78,12 @@ class ShapeValueGenerator(
     private fun classDeclaration(writer: KotlinWriter, shape: StructureShape, block: () -> Unit) {
         val symbol = symbolProvider.toSymbol(shape)
         // invoke the generated DSL builder for the class
-        writer.writeInline("#L {\n", symbol.name)
+        writer.writeInline("#L {", symbol.name)
+            .ensureNewline()
             .indent()
             .call { block() }
             .dedent()
-            .write("")
+            .ensureNewline()
             .write("}")
     }
 
@@ -72,11 +93,12 @@ class ShapeValueGenerator(
 
         val collectionGeneratorFunction = symbolProvider.toSymbol(shape).expectProperty(SymbolProperty.IMMUTABLE_COLLECTION_FUNCTION)
 
-        writer.writeInline("$collectionGeneratorFunction(\n")
+        writer.writeInline("$collectionGeneratorFunction(")
+            .ensureNewline()
             .indent()
             .call { block() }
             .dedent()
-            .write("")
+            .ensureNewline()
             .write(")")
 
         writer.popState()
@@ -92,11 +114,12 @@ class ShapeValueGenerator(
         collectionSymbol.references.forEach {
             writer.addImport(it.symbol)
         }
-        writer.writeInline("$generatorFn(\n")
+        writer.writeInline("$generatorFn(")
+            .ensureNewline()
             .indent()
             .call { block() }
             .dedent()
-            .write("")
+            .ensureNewline()
             .write(")")
 
         writer.popState()
@@ -142,7 +165,8 @@ class ShapeValueGenerator(
         override fun objectNode(node: ObjectNode) {
             if (currShape.type == ShapeType.DOCUMENT) {
                 writer
-                    .writeInline("#T {\n", RuntimeTypes.Core.Content.buildDocument)
+                    .writeInline("#T {", RuntimeTypes.Core.Content.buildDocument)
+                    .ensureNewline()
                     .indent()
             }
 
@@ -157,9 +181,9 @@ class ShapeValueGenerator(
                         memberShape = generator.model.expectShape(member.target)
                         val memberName = generator.symbolProvider.toMemberName(member)
                         writer.writeInline("#L = ", memberName)
-                        generator.writeShapeValueInline(writer, memberShape, valueNode)
+                        generator.instantiateShapeInline(writer, memberShape, valueNode)
                         if (i < node.members.size - 1) {
-                            writer.write("")
+                            writer.ensureNewline()
                         }
                     }
                     is MapShape -> {
@@ -169,16 +193,17 @@ class ShapeValueGenerator(
                         if (valueNode is NullNode) {
                             writer.write("null")
                         } else {
-                            generator.writeShapeValueInline(writer, memberShape, valueNode)
+                            generator.instantiateShapeInline(writer, memberShape, valueNode)
                             if (i < node.members.size - 1) {
-                                writer.writeInline(",\n")
+                                writer.writeInline(",")
+                                    .ensureNewline()
                             }
                         }
                     }
                     is DocumentShape -> {
                         writer.writeInline("#S to ", keyNode.value)
-                        generator.writeShapeValueInline(writer, currShape, valueNode)
-                        writer.writeInline("\n")
+                        generator.instantiateShapeInline(writer, currShape, valueNode)
+                        writer.ensureNewline()
                     }
                     is UnionShape -> {
                         val member = currShape.getMember(keyNode.value).orElseThrow {
@@ -189,8 +214,8 @@ class ShapeValueGenerator(
                         val memberName = generator.symbolProvider.toMemberName(member)
                         val variantName = memberName.replaceFirstChar { c -> c.uppercaseChar() }
                         writer.writeInline("${currSymbol.name}.$variantName(")
-                        generator.writeShapeValueInline(writer, memberShape, valueNode)
-                        writer.write(")")
+                        generator.instantiateShapeInline(writer, memberShape, valueNode)
+                        writer.writeInline(")")
                     }
                     else -> throw CodegenException("unexpected shape type " + currShape.type)
                 }
@@ -238,8 +263,10 @@ class ShapeValueGenerator(
                     writer.withInlineBlock("#T(", ")", RuntimeTypes.Core.Content.Document) {
                         writer.withInlineBlock("listOf(", ")") {
                             node.elements.forEach {
-                                generator.writeShapeValueInline(writer, currShape, it)
-                                writer.writeInline(",\n")
+                                generator.instantiateShapeInline(writer, currShape, it)
+                                writer.unwrite(writer.newline)
+                                writer.writeInline(",")
+                                    .ensureNewline()
                             }
                         }
                     }
@@ -249,9 +276,14 @@ class ShapeValueGenerator(
                     val memberShape = generator.model.expectShape((currShape as CollectionShape).member.target)
                     var i = 0
                     node.elements.forEach { element ->
-                        generator.writeShapeValueInline(writer, memberShape, element)
+                        generator.instantiateShapeInline(writer, memberShape, element)
+                        writer.unwrite(writer.newline)
                         if (i < node.elements.size - 1) {
-                            writer.writeInline(",\n")
+                            writer.pushState()
+                            writer.indentText = ""
+                            writer.writeInlineWithNoFormatting(",")
+                            writer.ensureNewline()
+                            writer.popState()
                         }
                         i++
                     }
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolUnitTestRequestGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolUnitTestRequestGenerator.kt
@@ -101,7 +101,7 @@ open class HttpProtocolUnitTestRequestGenerator protected constructor(builder: B
                     writer.writeInline("\nval input = ")
                         .indent()
                         .call {
-                            ShapeValueGenerator(model, symbolProvider).writeShapeValueInline(writer, inputShape, test.params)
+                            ShapeValueGenerator(model, symbolProvider).instantiateShapeInline(writer, inputShape, test.params)
                         }
                         .dedent()
                         .write("")
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolUnitTestResponseGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/protocol/HttpProtocolUnitTestResponseGenerator.kt
@@ -80,7 +80,7 @@ open class HttpProtocolUnitTestResponseGenerator protected constructor(builder:
                     .call {
                         outputShape?.let {
                             writer.writeInline("\nresponse = ")
-                            ShapeValueGenerator(model, symbolProvider).writeShapeValueInline(writer, it, test.params)
+                            ShapeValueGenerator(model, symbolProvider).instantiateShapeInline(writer, it, test.params)
                         }
                     }
                     .write("")
diff --git a/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/samples/KDocSamplesGenerator.kt b/codegen/smithy-kotlin-codegen/src/main/kotlin/software/amazon/smithy/kotlin/codegen/rendering/samples/KDocSamplesGenerator.kt
@@ -0,0 +1,157 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+package software.amazon.smithy.kotlin.codegen.rendering.samples
+
+import software.amazon.smithy.kotlin.codegen.KotlinSettings
+import software.amazon.smithy.kotlin.codegen.core.*
+import software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
+import software.amazon.smithy.kotlin.codegen.model.expectTrait
+import software.amazon.smithy.kotlin.codegen.model.getTrait
+import software.amazon.smithy.kotlin.codegen.model.hasAllOptionalMembers
+import software.amazon.smithy.kotlin.codegen.model.hasTrait
+import software.amazon.smithy.kotlin.codegen.rendering.ShapeValueGenerator
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.model.SourceLocation
+import software.amazon.smithy.model.knowledge.TopDownIndex
+import software.amazon.smithy.model.shapes.OperationShape
+import software.amazon.smithy.model.traits.DocumentationTrait
+import software.amazon.smithy.model.traits.ExamplesTrait
+import software.amazon.smithy.model.traits.ExamplesTrait.Example
+import software.amazon.smithy.model.transform.ModelTransformer
+import java.util.*
+import kotlin.jvm.optionals.getOrDefault
+
+/**
+ * [KotlinIntegration] that renders [KDoc samples](https://kotlinlang.org/docs/kotlin-doc.html#sample-identifier)
+ * and pre-processes the documentation to insert references to the generated sample identifiers for operations
+ * that have [examples](https://smithy.io/2.0/spec/documentation-traits.html#smithy-api-examples-trait)
+ */
+class KDocSamplesGenerator : KotlinIntegration {
+    override fun enabledForService(model: Model, settings: KotlinSettings): Boolean {
+        val topDownIndex = TopDownIndex.of(model)
+        val operations = topDownIndex.getContainedOperations(settings.service)
+        return operations.any { it.hasTrait<ExamplesTrait>() }
+    }
+
+    /**
+     * This should run _after_ [software.amazon.smithy.kotlin.codegen.lang.DocumentationPreprocessor]
+     */
+    override fun preprocessModel(model: Model, settings: KotlinSettings): Model {
+        val transformer = ModelTransformer.create()
+        return transformer.mapShapes(model) { shape ->
+            when {
+                shape is OperationShape && shape.hasTrait<ExamplesTrait>() -> {
+                    val examplesTrait = shape.expectTrait<ExamplesTrait>()
+                    val filtered = examplesTrait.examples.filterNot { it.error.isPresent }
+                    val kdocSampleIdentifiers = filtered.indices.joinToString(separator = "\n") { idx ->
+                        val identifier = sampleIdentifier(settings, shape, idx)
+                        "@sample $identifier"
+                    }
+
+                    val existingDocs = shape.getTrait<DocumentationTrait>()
+                    val updatedDocs = buildString {
+                        if (existingDocs != null) {
+                            append(existingDocs.value)
+                            append("\n\n")
+                        }
+                        append(kdocSampleIdentifiers)
+                    }
+                    val sourceLocation = existingDocs?.sourceLocation ?: SourceLocation.NONE
+                    val newOrUpdatedDocTrait = DocumentationTrait(updatedDocs, sourceLocation)
+                    shape.toBuilder()
+                        .addTrait(newOrUpdatedDocTrait)
+                        .build()
+                }
+                else -> shape
+            }
+        }
+    }
+
+    private fun sampleIdentifier(settings: KotlinSettings, op: OperationShape, index: Int): String =
+        listOf(
+            samplePackage(settings),
+            sampleClassName(op),
+            sampleFunctionName(index),
+        ).joinToString(separator = ".")
+
+    private fun sampleFunctionName(index: Int): String = "sample" + if (index > 0) "${index + 1}" else ""
+    private fun sampleClassName(op: OperationShape): String = op.id.name
+
+    private fun samplePackage(settings: KotlinSettings): String = settings.pkg.subpackage("samples")
+
+    override fun writeAdditionalFiles(ctx: CodegenContext, delegator: KotlinDelegator) {
+        val topDownIndex = TopDownIndex.of(ctx.model)
+        val operations = topDownIndex.getContainedOperations(ctx.settings.service)
+        operations.filter { it.hasTrait<ExamplesTrait>() }
+            .forEach { op ->
+                val examples = op.expectTrait<ExamplesTrait>()
+
+                val writer = KotlinWriter(samplePackage(ctx.settings))
+                writer.withBlock("class #L {", "}", sampleClassName(op)) {
+                    examples
+                        .examples
+                        // unclear what benefit error examples provide, omit for now
+                        .filterNot { it.error.isPresent }
+                        .forEachIndexed { idx, example ->
+                            write("")
+                            write("@Sample")
+                            withBlock("fun #L() {", "}", sampleFunctionName(idx)) {
+                                example
+                                    .documentation
+                                    .getOrDefault(example.title)
+                                    .breakLongLines()
+                                    .forEach { line ->
+                                        write("// #L", line)
+                                    }
+
+                                renderNormalExample(ctx, writer, op, example)
+                            }
+                        }
+                }
+                val contents = writer.toString()
+                delegator.fileManifest.writeFile("src/samples/${op.id.name}.kt", contents)
+            }
+    }
+
+    private fun renderNormalExample(ctx: CodegenContext, writer: KotlinWriter, op: OperationShape, example: Example) {
+        val clientName = clientName(ctx.settings.sdkId).replaceFirstChar { it.lowercase(Locale.getDefault()) }
+        val respPrefix = if (example.output.isPresent) "val resp = " else ""
+
+        val input = ctx.model.expectShape(op.inputShape)
+        if (input.hasAllOptionalMembers && example.input.isEmpty) {
+            writer.write("#L#LClient.#L()", respPrefix, clientName, op.defaultName())
+        } else {
+            writer.withBlock("#L#LClient.#L {", "}", respPrefix, clientName, op.defaultName()) {
+                ShapeValueGenerator(ctx.model, ctx.symbolProvider).writeShapeValues(writer, input, example.input)
+            }
+        }
+    }
+}
+
+private val wordsPattern = Regex("""\w+[.,]?|".*?"[.,]?|\(.*\)[.,]?""")
+internal fun String.breakLongLines(maxLineLengthChars: Int = 100): List<String> {
+    val words = wordsPattern.findAll(this).map(MatchResult::value)
+    val lines = mutableListOf<String>()
+    val wordsOnLine = mutableListOf<String>()
+    var lineLength = 0
+
+    words.forEach { word ->
+        if (word.length + lineLength < maxLineLengthChars) {
+            if (wordsOnLine.isNotEmpty()) lineLength++
+            lineLength += word.length
+            wordsOnLine.add(word)
+        } else {
+            lines.add(wordsOnLine.joinToString(separator = " "))
+            lineLength = 0
+            wordsOnLine.clear()
+            wordsOnLine.add(word)
+        }
+    }
+
+    if (wordsOnLine.isNotEmpty()) {
+        lines.add(wordsOnLine.joinToString(separator = " "))
+    }
+    return lines
+}
diff --git a/codegen/smithy-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration b/codegen/smithy-kotlin-codegen/src/main/resources/META-INF/services/software.amazon.smithy.kotlin.codegen.integration.KotlinIntegration
@@ -1,5 +1,6 @@
 software.amazon.smithy.kotlin.codegen.lang.BuiltinPreprocessor
 software.amazon.smithy.kotlin.codegen.lang.DocumentationPreprocessor
+software.amazon.smithy.kotlin.codegen.rendering.samples.KDocSamplesGenerator
 software.amazon.smithy.kotlin.codegen.model.SetRefactorPreprocessor
 software.amazon.smithy.kotlin.codegen.rendering.PaginatorGenerator
 software.amazon.smithy.kotlin.codegen.rendering.waiters.ServiceWaitersGenerator
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ShapeValueGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/ShapeValueGeneratorTest.kt
diff --git a/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/samples/KDocSamplesGeneratorTest.kt b/codegen/smithy-kotlin-codegen/src/test/kotlin/software/amazon/smithy/kotlin/codegen/rendering/samples/KDocSamplesGeneratorTest.kt
diff --git a/docs/dokka-presets/css/aws-styles.css b/docs/dokka-presets/css/aws-styles.css

Original file line number	Diff line number	Diff line change
`@@ -101,7 +101,7 @@ open class HttpProtocolUnitTestRequestGenerator protected constructor(builder: B`
`101`	`101`	`writer.writeInline("\nval input = ")`
`102`	`102`	`.indent()`
`103`	`103`	`.call {`
`104`		`- ShapeValueGenerator(model, symbolProvider).writeShapeValueInline(writer, inputShape, test.params)`
	`104`	`+ ShapeValueGenerator(model, symbolProvider).instantiateShapeInline(writer, inputShape, test.params)`
`105`	`105`	`}`
`106`	`106`	`.dedent()`
`107`	`107`	`.write("")`
Original file line number	Diff line number	Diff line change
`@@ -80,7 +80,7 @@ open class HttpProtocolUnitTestResponseGenerator protected constructor(builder:`
`80`	`80`	`.call {`
`81`	`81`	`outputShape?.let {`
`82`	`82`	`writer.writeInline("\nresponse = ")`
`83`		`- ShapeValueGenerator(model, symbolProvider).writeShapeValueInline(writer, it, test.params)`
	`83`	`+ ShapeValueGenerator(model, symbolProvider).instantiateShapeInline(writer, it, test.params)`
`84`	`84`	`}`
`85`	`85`	`}`
`86`	`86`	`.write("")`