Unverified Commit 1abb0180 authored by John DiSanti's avatar John DiSanti Committed by GitHub
Browse files

Introduce syntax sugar for accessing Smithy shape traits (#385)

parent 743d3e89

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/AwsEndpointDecorator.kt

+2 −1
Original line number Diff line number Diff line
@@ -23,6 +23,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig 
import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustomization

import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig

import software.amazon.smithy.rust.codegen.util.dq

import software.amazon.smithy.rust.codegen.util.expectTrait



class AwsEndpointDecorator : RustCodegenDecorator {

    override val name: String = "AwsEndpoint"
@@ -56,7 +57,7 @@ class AwsEndpointDecorator : RustCodegenDecorator { 


class EndpointConfigCustomization(private val runtimeConfig: RuntimeConfig, serviceShape: ServiceShape) :

    ConfigCustomization() {

    private val endpointPrefix = serviceShape.expectTrait(ServiceTrait::class.java).endpointPrefix

    private val endpointPrefix = serviceShape.expectTrait<ServiceTrait>().endpointPrefix

    private val resolveAwsEndpoint = runtimeConfig.awsEndpointDependency().asType().copy(name = "ResolveAwsEndpoint")

    override fun section(section: ServiceConfig): Writable = writable {

        when (section) {

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/SigV4SigningDecorator.kt

+4 −2
Original line number Diff line number Diff line
@@ -21,6 +21,8 @@ import software.amazon.smithy.rust.codegen.smithy.generators.config.ConfigCustom 
import software.amazon.smithy.rust.codegen.smithy.generators.config.ServiceConfig

import software.amazon.smithy.rust.codegen.smithy.letIf

import software.amazon.smithy.rust.codegen.util.dq

import software.amazon.smithy.rust.codegen.util.expectTrait

import software.amazon.smithy.rust.codegen.util.hasTrait



/**

 * The SigV4SigningDecorator:
@@ -33,14 +35,14 @@ class SigV4SigningDecorator : RustCodegenDecorator { 
    override val name: String = "SigV4Signing"

    override val order: Byte = 0



    private fun applies(protocolConfig: ProtocolConfig): Boolean = protocolConfig.serviceShape.hasTrait(SigV4Trait::class.java)

    private fun applies(protocolConfig: ProtocolConfig): Boolean = protocolConfig.serviceShape.hasTrait<SigV4Trait>()



    override fun configCustomizations(

        protocolConfig: ProtocolConfig,

        baseCustomizations: List<ConfigCustomization>

    ): List<ConfigCustomization> {

        return baseCustomizations.letIf(applies(protocolConfig)) {

            it + SigV4SigningConfig(protocolConfig.serviceShape.expectTrait(SigV4Trait::class.java))

            it + SigV4SigningConfig(protocolConfig.serviceShape.expectTrait())

        }

    }

aws/sdk-codegen/src/main/kotlin/software/amazon/smithy/rustsdk/UserAgentDecorator.kt

+2 −1
Original line number Diff line number Diff line
@@ -20,6 +20,7 @@ import software.amazon.smithy.rust.codegen.smithy.generators.LibRsCustomization 
import software.amazon.smithy.rust.codegen.smithy.generators.LibRsSection

import software.amazon.smithy.rust.codegen.smithy.generators.ProtocolConfig

import software.amazon.smithy.rust.codegen.util.dq

import software.amazon.smithy.rust.codegen.util.expectTrait



/**

 * Inserts a UserAgent configuration into the operation
@@ -33,7 +34,7 @@ class UserAgentDecorator : RustCodegenDecorator { 
        baseCustomizations: List<LibRsCustomization>

    ): List<LibRsCustomization> {

        // We are generating an AWS SDK, the service needs to have the AWS service trait

        val serviceTrait = protocolConfig.serviceShape.expectTrait(ServiceTrait::class.java)

        val serviceTrait = protocolConfig.serviceShape.expectTrait<ServiceTrait>()

        return baseCustomizations + ApiVersion(protocolConfig.runtimeConfig, serviceTrait)

    }

aws/sdk-codegen/src/test/kotlin/software/amazon/smithy/rustsdk/EndpointConfigCustomizationTest.kt

+2 −1
Original line number Diff line number Diff line
@@ -15,6 +15,7 @@ import software.amazon.smithy.rust.codegen.testutil.compileAndTest 
import software.amazon.smithy.rust.codegen.testutil.stubConfigProject

import software.amazon.smithy.rust.codegen.testutil.unitTest

import software.amazon.smithy.rust.codegen.testutil.validateConfigCustomizations

import software.amazon.smithy.rust.codegen.util.expectTrait

import software.amazon.smithy.rust.codegen.util.lookup



internal class EndpointConfigCustomizationTest {
@@ -41,7 +42,7 @@ internal class EndpointConfigCustomizationTest { 
    fun `generates valid code when no endpoint prefix is provided`() {

        val serviceShape = model.lookup<ServiceShape>("test#NoEndpointPrefix")

        validateConfigCustomizations(EndpointConfigCustomization(AwsTestRuntimeConfig, serviceShape))

        serviceShape.expectTrait(ServiceTrait::class.java).endpointPrefix shouldBe "noendpointprefix"

        serviceShape.expectTrait<ServiceTrait>().endpointPrefix shouldBe "noendpointprefix"

    }



    @Test

codegen/src/main/kotlin/software/amazon/smithy/rust/codegen/smithy/CodegenVisitor.kt

+4 −2
Original line number Diff line number Diff line
@@ -29,6 +29,8 @@ import software.amazon.smithy.rust.codegen.smithy.protocols.ProtocolLoader 
import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticInputTrait

import software.amazon.smithy.rust.codegen.smithy.transformers.RecursiveShapeBoxer

import software.amazon.smithy.rust.codegen.util.CommandFailed

import software.amazon.smithy.rust.codegen.util.getTrait

import software.amazon.smithy.rust.codegen.util.hasTrait

import software.amazon.smithy.rust.codegen.util.runCommand

import java.util.logging.Logger
@@ -103,7 +105,7 @@ class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustC 
        logger.fine("generating a structure...")

        rustCrate.useShapeWriter(shape) { writer ->

            StructureGenerator(model, symbolProvider, writer, shape).render()

            if (!shape.hasTrait(SyntheticInputTrait::class.java)) {

            if (!shape.hasTrait<SyntheticInputTrait>()) {

                val builderGenerator = BuilderGenerator(protocolConfig.model, protocolConfig.symbolProvider, shape)

                builderGenerator.render(writer)

                writer.implBlock(shape, symbolProvider) {
@@ -114,7 +116,7 @@ class CodegenVisitor(context: PluginContext, private val codegenDecorator: RustC 
    }



    override fun stringShape(shape: StringShape) {

        shape.getTrait(EnumTrait::class.java).map { enum ->

        shape.getTrait<EnumTrait>()?.also { enum ->

            rustCrate.useShapeWriter(shape) { writer ->

                EnumGenerator(symbolProvider, writer, shape, enum).render()

            }