Fix bug in coroutine scheduling. Add ByteStream implementation and simplify Python symbol visitor.

codegen-server-test/python/model/pokemon.smithy

@@ -13,7 +13,7 @@ use aws.protocols#restJson1
 service PokemonService {
     version: "2021-12-01",
     resources: [PokemonSpecies],
-    operations: [GetServerStatistics, EmptyOperation, HealthCheckOperation],
+    operations: [GetServerStatistics, EmptyOperation, HealthCheckOperation, StreamPokemonRadioOperation],
 }
 
 /// A Pokémon species forms the basis for at least one Pokémon.
@@ -70,6 +70,22 @@ structure GetServerStatisticsOutput {
     calls_count: Long,
 }
 
+/// Fetch the radio song from the database and stream it back as a playable audio.
+@readonly
+@http(uri: "/radio", method: "GET")
+operation StreamPokemonRadioOperation {
+    output: StreamPokemonRadioOutput,
+}
+
+@output
+structure StreamPokemonRadioOutput {
+    @httpPayload
+    data: StreamingBlob,
+}
+
+@streaming
+blob StreamingBlob
+
 list FlavorTextEntries {
     member: FlavorText
 }
@@ -127,8 +143,7 @@ structure EmptyOperationOutput { }
 /// Not yet a deep check
 @readonly
 @http(uri: "/ping", method: "GET")
-operation HealthCheckOperation {
-}
+operation HealthCheckOperation { }
 
 @error("client")
 @httpError(404)

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonCodegenServerPlugin.kt

@@ -12,13 +12,12 @@ import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.rust.codegen.rustlang.RustReservedWordSymbolProvider
 import software.amazon.smithy.rust.codegen.server.python.smithy.customizations.DECORATORS
-import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerSymbolProvider
+import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonServerSymbolVisitor
+import software.amazon.smithy.rust.codegen.server.python.smithy.generators.PythonStreamingShapeMetadataProvider
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
 import software.amazon.smithy.rust.codegen.smithy.BaseSymbolMetadataProvider
 import software.amazon.smithy.rust.codegen.smithy.EventStreamSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.ServerCodegenContext
-import software.amazon.smithy.rust.codegen.smithy.StreamingShapeMetadataProvider
-import software.amazon.smithy.rust.codegen.smithy.StreamingShapeSymbolProvider
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitor
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitorConfig
 import software.amazon.smithy.rust.codegen.smithy.customize.CombinedCodegenDecorator
@@ -68,20 +67,17 @@ class PythonCodegenServerPlugin : SmithyBuildPlugin {
             serviceShape: ServiceShape,
             symbolVisitorConfig: SymbolVisitorConfig,
         ) =
-            SymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
+            // Rename a set of symbols that do not implement `PyClass` and have been wrapped in
+            // `aws_smithy_http_server_python::types`.
+            PythonServerSymbolVisitor(model, serviceShape = serviceShape, config = symbolVisitorConfig)
                 // Generate different types for EventStream shapes (e.g. transcribe streaming)
                 .let {
                     EventStreamSymbolProvider(symbolVisitorConfig.runtimeConfig, it, model, CodegenTarget.SERVER)
                 }
-                // Generate [ByteStream] instead of `Blob` for streaming binary shapes (e.g. S3 GetObject)
-                .let { StreamingShapeSymbolProvider(it, model) }
-                // Rename a set of symbols that do not implement `PyClass` and have been wrapped in
-                // `aws_smithy_http_server_python::types`.
-                .let { PythonServerSymbolProvider(it, model) }
                 // Add Rust attributes (like `#[derive(PartialEq)]`) to generated shapes
                 .let { BaseSymbolMetadataProvider(it, model, additionalAttributes = listOf()) }
                 // Streaming shapes need different derives (e.g. they cannot derive Eq)
-                .let { StreamingShapeMetadataProvider(it, model) }
+                .let { PythonStreamingShapeMetadataProvider(it, model) }
                 // Rename shapes that clash with Rust reserved words & and other SDK specific features e.g. `send()` cannot
                 // be the name of an operation input
                 .let { RustReservedWordSymbolProvider(it, model) }

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCargoDependency.kt

@@ -23,6 +23,7 @@ object PythonServerCargoDependency {
     val TowerHttp: CargoDependency = CargoDependency("tower-http", CratesIo("0.3"), features = setOf("trace"))
     val Hyper: CargoDependency = CargoDependency("hyper", CratesIo("0.14.12"), features = setOf("server", "http1", "http2", "tcp", "stream"))
     val NumCpus: CargoDependency = CargoDependency("num_cpus", CratesIo("1.13"))
+    val ParkingLot: CargoDependency = CargoDependency("parking_lot", CratesIo("0.12"))
 
     fun SmithyHttpServer(runtimeConfig: RuntimeConfig) = runtimeConfig.runtimeCrate("http-server")
     fun SmithyHttpServerPython(runtimeConfig: RuntimeConfig) = runtimeConfig.runtimeCrate("http-server-python")

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCodegenVisitor.kt

@@ -27,7 +27,6 @@ import software.amazon.smithy.rust.codegen.smithy.generators.BuilderGenerator
 import software.amazon.smithy.rust.codegen.smithy.generators.CodegenTarget
 import software.amazon.smithy.rust.codegen.smithy.generators.implBlock
 import software.amazon.smithy.rust.codegen.util.getTrait
-import software.amazon.smithy.rust.codegen.util.hasStreamingMember
 
 /**
  * Entrypoint for Python server-side code generation. This class will walk the in-memory model and
@@ -83,9 +82,6 @@ class PythonServerCodegenVisitor(
      * This function _does not_ generate any serializers.
      */
     override fun structureShape(shape: StructureShape) {
-        if (shape.hasStreamingMember(model)) {
-            throw CodegenException("Streaming members are not supported in Python yet")
-        }
         logger.info("[python-server-codegen] Generating a structure $shape")
         rustCrate.useShapeWriter(shape) { writer ->
             // Use Python specific structure generator that adds the #[pyclass] attribute

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerRuntimeType.kt

@@ -21,6 +21,9 @@ object PythonServerRuntimeType {
     fun Blob(runtimeConfig: RuntimeConfig) =
         RuntimeType("Blob", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
 
+    fun ByteStream(runtimeConfig: RuntimeConfig) =
+        RuntimeType("ByteStream", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
+
     fun DateTime(runtimeConfig: RuntimeConfig) =
         RuntimeType("DateTime", PythonServerCargoDependency.SmithyHttpServerPython(runtimeConfig), "${runtimeConfig.crateSrcPrefix}_http_server_python::types")
 

codegen-server/python/src/main/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerSymbolProvider.kt

@@ -8,102 +8,107 @@ package software.amazon.smithy.rust.codegen.server.python.smithy.generators
 import software.amazon.smithy.codegen.core.Symbol
 import software.amazon.smithy.model.Model
 import software.amazon.smithy.model.shapes.BlobShape
-import software.amazon.smithy.model.shapes.ListShape
-import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
-import software.amazon.smithy.model.shapes.SetShape
+import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
-import software.amazon.smithy.model.shapes.ShapeVisitor
+import software.amazon.smithy.model.shapes.StringShape
 import software.amazon.smithy.model.shapes.StructureShape
 import software.amazon.smithy.model.shapes.TimestampShape
 import software.amazon.smithy.model.shapes.UnionShape
-import software.amazon.smithy.model.traits.EnumDefinition
-import software.amazon.smithy.rust.codegen.rustlang.RustType
+import software.amazon.smithy.rust.codegen.rustlang.RustMetadata
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerRuntimeType
-import software.amazon.smithy.rust.codegen.smithy.MaybeRenamed
+import software.amazon.smithy.rust.codegen.smithy.RuntimeType
 import software.amazon.smithy.rust.codegen.smithy.RustSymbolProvider
+import software.amazon.smithy.rust.codegen.smithy.SymbolMetadataProvider
+import software.amazon.smithy.rust.codegen.smithy.SymbolVisitor
 import software.amazon.smithy.rust.codegen.smithy.SymbolVisitorConfig
-import software.amazon.smithy.rust.codegen.smithy.rustType
+import software.amazon.smithy.rust.codegen.smithy.expectRustMetadata
 import software.amazon.smithy.rust.codegen.smithy.shape
-import software.amazon.smithy.rust.codegen.util.toPascalCase
-import software.amazon.smithy.rust.codegen.util.toSnakeCase
+import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticInputTrait
+import software.amazon.smithy.rust.codegen.smithy.traits.SyntheticOutputTrait
+import software.amazon.smithy.rust.codegen.util.hasStreamingMember
+import software.amazon.smithy.rust.codegen.util.hasTrait
+import software.amazon.smithy.rust.codegen.util.isStreaming
 
-// Wrapping symbol visitor provider allowing to implement symbol providers that can recursively replace symbols in nested shapes.
-open class PythonWrappingVisitingSymbolProvider(private val base: RustSymbolProvider, private val model: Model) : ShapeVisitor.Default<Symbol>(), RustSymbolProvider {
-    override fun getDefault(shape: Shape): Symbol {
-        return base.toSymbol(shape)
-    }
-
-    override fun config(): SymbolVisitorConfig {
-        return base.config()
-    }
-
-    override fun toEnumVariantName(definition: EnumDefinition): MaybeRenamed? {
-        return base.toEnumVariantName(definition)
-    }
+/**
+ * Symbol visitor  allowing that recursively replace symbols in nested shapes.
+ *
+ * Input / output / error structures can refer to complex types like the ones implemented inside
+ * `aws_smithy_types` (a good example is `aws_smithy_types::Blob`).
+ * `aws_smithy_http_server_python::types` wraps those types that do not implement directly the
+ * `pyo3::PyClass` trait and cannot be shared safely with Python, providing an idiomatic Python / Rust API.
+ *
+ * This symbol provider ensures types not implementing `pyo3::PyClass` are swapped with their wrappers from
+ * `aws_smithy_http_server_python::types`.
+ */
+class PythonServerSymbolVisitor(
+    private val model: Model,
+    private val serviceShape: ServiceShape?,
+    private val config: SymbolVisitorConfig,
+) : SymbolVisitor(model, serviceShape, config) {
+    private val runtimeConfig = config().runtimeConfig
 
-    override fun toMemberName(shape: MemberShape): String = when (val container = model.expectShape(shape.container)) {
-        is StructureShape -> shape.memberName.toSnakeCase()
-        is UnionShape -> shape.memberName.toPascalCase()
-        else -> error("unexpected container shape: $container")
-    }
+    override fun toSymbol(shape: Shape): Symbol {
+        val initial = shape.accept(this)
 
-    override fun listShape(shape: ListShape): Symbol {
-        val inner = toSymbol(shape.member)
-        return symbolBuilder(RustType.Vec(inner.rustType())).addReference(inner).build()
-    }
+        if (shape !is MemberShape) {
+            return initial
+        }
+        val target = model.expectShape(shape.target)
+        val container = model.expectShape(shape.container)
 
-    override fun mapShape(shape: MapShape): Symbol {
-        val target = model.expectShape(shape.key.target)
-        require(target.isStringShape) { "unexpected key shape: ${shape.key}: $target [keys must be strings]" }
-        val key = toSymbol(shape.key)
-        val value = toSymbol(shape.value)
-        return symbolBuilder(RustType.HashMap(key.rustType(), value.rustType())).addReference(key)
-            .addReference(value).build()
-    }
+        // We are only targetting non syntetic inputs and outputs.
+        if (!container.hasTrait<SyntheticOutputTrait>() && !container.hasTrait<SyntheticInputTrait>()) {
+            return initial
+        }
 
-    override fun setShape(shape: SetShape): Symbol {
-        val inner = toSymbol(shape.member)
-        val builder = if (model.expectShape(shape.member.target).isStringShape) {
-            symbolBuilder(RustType.HashSet(inner.rustType()))
+        // We are only targeting streaming blobs as the rest of the symbols do not change if streaming is enabled.
+        // For example a TimestampShape doesn't become a different symbol when streaming is involved, but BlobShape
+        // become a ByteStream.
+        return if (target is BlobShape && shape.isStreaming(model)) {
+            PythonServerRuntimeType.ByteStream(config().runtimeConfig).toSymbol()
         } else {
-            // only strings get put into actual sets because floats are unhashable
-            symbolBuilder(RustType.Vec(inner.rustType()))
+            initial
         }
-        return builder.addReference(inner).build()
     }
 
-    override fun memberShape(shape: MemberShape): Symbol {
-        return toSymbol(model.expectShape(shape.target))
-    }
-
-    override fun toSymbol(shape: Shape): Symbol {
-        return shape.accept(this)
+    override fun timestampShape(shape: TimestampShape?): Symbol {
+        return PythonServerRuntimeType.DateTime(runtimeConfig).toSymbol()
     }
 
-    private fun symbolBuilder(rustType: RustType): Symbol.Builder {
-        return Symbol.builder().rustType(rustType)
-            .definitionFile("python.rs")
+    override fun blobShape(shape: BlobShape?): Symbol {
+        return PythonServerRuntimeType.Blob(runtimeConfig).toSymbol()
     }
 }
 
 /**
- * Input / output / error structures can refer to complex types like the ones implemented inside
- * `aws_smithy_types` (a good example is `aws_smithy_types::Blob`).
- * `aws_smithy_http_server_python::types` wraps those types that do not implement directly the
- * `pyo3::PyClass` trait and cannot be shared safely with Python, providing an idiomatic Python / Rust API.
+ * SymbolProvider to drop the PartialEq bounds in streaming shapes
  *
- * This symbol provider ensures types not implementing `pyo3::PyClass` are swapped with their wrappers from
- * `aws_smithy_http_server_python::types`.
+ * Streaming shapes equality cannot be checked without reading the body. Because of this, these shapes
+ * do not implement `PartialEq`.
+ *
+ * Note that since streaming members can only be used on the root shape, this can only impact input and output shapes.
  */
-class PythonServerSymbolProvider(private val base: RustSymbolProvider, private val model: Model) : PythonWrappingVisitingSymbolProvider(base, model) {
-    private val runtimeConfig = config().runtimeConfig
+class PythonStreamingShapeMetadataProvider(private val base: RustSymbolProvider, private val model: Model) : SymbolMetadataProvider(base) {
+    override fun memberMeta(memberShape: MemberShape): RustMetadata {
+        return base.toSymbol(memberShape).expectRustMetadata()
+    }
 
-    override fun timestampShape(shape: TimestampShape?): Symbol {
-        return PythonServerRuntimeType.DateTime(runtimeConfig).toSymbol()
+    override fun structureMeta(structureShape: StructureShape): RustMetadata {
+        val baseMetadata = base.toSymbol(structureShape).expectRustMetadata()
+        return if (structureShape.hasStreamingMember(model)) {
+            baseMetadata.withoutDerives(RuntimeType.PartialEq)
+        } else baseMetadata
     }
 
-    override fun blobShape(shape: BlobShape?): Symbol {
-        return PythonServerRuntimeType.Blob(runtimeConfig).toSymbol()
+    override fun unionMeta(unionShape: UnionShape): RustMetadata {
+        val baseMetadata = base.toSymbol(unionShape).expectRustMetadata()
+        return if (unionShape.hasStreamingMember(model)) {
+            baseMetadata.withoutDerives(RuntimeType.PartialEq)
+        } else baseMetadata
+    }
+
+    override fun enumMeta(stringShape: StringShape): RustMetadata {
+        return base.toSymbol(stringShape).expectRustMetadata()
     }
 }