Python: Support document type

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

@@ -15,7 +15,7 @@ import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
  * For a dependency that is used in the client, or in both the client and the server, use [CargoDependency] directly.
  */
 object PythonServerCargoDependency {
-    val PyO3: CargoDependency = CargoDependency("pyo3", CratesIo("0.17"), features = setOf("extension-module"))
+    val PyO3: CargoDependency = CargoDependency("pyo3", CratesIo("0.17"))
     val PyO3Asyncio: CargoDependency = CargoDependency("pyo3-asyncio", CratesIo("0.17"), features = setOf("attributes", "tokio-runtime"))
     val Tokio: CargoDependency = CargoDependency("tokio", CratesIo("1.20.1"), features = setOf("full"))
     val Tracing: CargoDependency = CargoDependency("tracing", CratesIo("0.1"))

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

@@ -22,4 +22,7 @@ object PythonServerRuntimeType {
 
     fun dateTime(runtimeConfig: RuntimeConfig) =
         PythonServerCargoDependency.smithyHttpServerPython(runtimeConfig).toType().resolve("types::DateTime")
+
+    fun document(runtimeConfig: RuntimeConfig) =
+        PythonServerCargoDependency.smithyHttpServerPython(runtimeConfig).toType().resolve("types::Document")
 }

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

@@ -8,6 +8,7 @@ package software.amazon.smithy.rust.codegen.server.python.smithy
 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.DocumentShape
 import software.amazon.smithy.model.shapes.ListShape
 import software.amazon.smithy.model.shapes.MapShape
 import software.amazon.smithy.model.shapes.MemberShape
@@ -80,6 +81,10 @@ class PythonServerSymbolVisitor(
     override fun blobShape(shape: BlobShape?): Symbol {
         return PythonServerRuntimeType.blob(runtimeConfig).toSymbol()
     }
+
+    override fun documentShape(shape: DocumentShape?): Symbol {
+        return PythonServerRuntimeType.document(runtimeConfig).toSymbol()
+    }
 }
 
 /**

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

@@ -6,6 +6,7 @@
 package software.amazon.smithy.rust.codegen.server.python.smithy.customizations
 
 import com.moandjiezana.toml.TomlWriter
+import software.amazon.smithy.rust.codegen.core.rustlang.Feature
 import software.amazon.smithy.rust.codegen.core.rustlang.Writable
 import software.amazon.smithy.rust.codegen.core.rustlang.docs
 import software.amazon.smithy.rust.codegen.core.rustlang.rust
@@ -57,6 +58,7 @@ class PubUsePythonTypes(private val codegenContext: ServerCodegenContext) : LibR
                 rustBlock("pub mod python_types") {
                     rust("pub use #T;", PythonServerRuntimeType.blob(codegenContext.runtimeConfig).toSymbol())
                     rust("pub use #T;", PythonServerRuntimeType.dateTime(codegenContext.runtimeConfig).toSymbol())
+                    rust("pub use #T;", PythonServerRuntimeType.document(codegenContext.runtimeConfig).toSymbol())
                 }
             }
             else -> emptySection
@@ -114,6 +116,24 @@ class PyProjectTomlDecorator : ServerCodegenDecorator {
     }
 }
 
+/**
+ * Adds `pyo3/extension-module` feature to default features.
+ *
+ * To be able to run `cargo test` with PyO3 we need two things:
+ *  - Make `pyo3/extension-module` optional and default
+ *  - Run tests with `cargo test --no-default-features`
+ * See: https://pyo3.rs/main/faq#i-cant-run-cargo-test-or-i-cant-build-in-a-cargo-workspace-im-having-linker-issues-like-symbol-not-found-or-undefined-reference-to-_pyexc_systemerror
+ */
+class PyO3ExtensionModuleDecorator : ServerCodegenDecorator {
+    override val name: String = "PyO3ExtensionModuleDecorator"
+    override val order: Byte = 0
+
+    override fun extras(codegenContext: ServerCodegenContext, rustCrate: RustCrate) {
+        // Add `pyo3/extension-module` to default features.
+        rustCrate.mergeFeature(Feature("extension-module", true, listOf("pyo3/extension-module")))
+    }
+}
+
 val DECORATORS = listOf(
     /**
      * Add the [InternalServerError] error to all operations.
@@ -128,4 +148,6 @@ val DECORATORS = listOf(
     PythonExportModuleDecorator(),
     // Generate `pyproject.toml` for the crate.
     PyProjectTomlDecorator(),
+    // Add PyO3 extension module feature.
+    PyO3ExtensionModuleDecorator(),
 )

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

@@ -0,0 +1,46 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.python.smithy.testutil
+
+import software.amazon.smithy.build.PluginContext
+import software.amazon.smithy.model.Model
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeConfig
+import software.amazon.smithy.rust.codegen.core.smithy.RuntimeCrateLocation
+import software.amazon.smithy.rust.codegen.core.testutil.generatePluginContext
+import software.amazon.smithy.rust.codegen.core.util.runCommand
+import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerCodegenVisitor
+import software.amazon.smithy.rust.codegen.server.python.smithy.customizations.DECORATORS
+import software.amazon.smithy.rust.codegen.server.smithy.customizations.ServerRequiredCustomizations
+import software.amazon.smithy.rust.codegen.server.smithy.customize.CombinedServerCodegenDecorator
+import java.io.File
+import java.nio.file.Path
+
+val TestRuntimeConfig =
+    RuntimeConfig(runtimeCrateLocation = RuntimeCrateLocation.Path(File("../../rust-runtime").absolutePath))
+
+fun generatePythonServerPluginContext(model: Model) =
+    generatePluginContext(model, runtimeConfig = TestRuntimeConfig)
+
+fun executePythonServerCodegenVisitor(pluginCtx: PluginContext) {
+    val codegenDecorator: CombinedServerCodegenDecorator =
+        CombinedServerCodegenDecorator.fromClasspath(
+            pluginCtx,
+            CombinedServerCodegenDecorator(DECORATORS + ServerRequiredCustomizations()),
+        )
+    PythonServerCodegenVisitor(pluginCtx, codegenDecorator).execute()
+}
+
+fun cargoTest(workdir: Path) =
+    // `--no-default-features` is required to disable `pyo3/extension-module` which causes linking errors
+    // see `PyO3ExtensionModuleDecorator`'s comments fore more detail.
+    "cargo test --no-default-features".runCommand(
+        workdir,
+        mapOf(
+            // Those are required to run tests on macOS, see: https://pyo3.rs/main/building_and_distribution#macos
+            "CARGO_TARGET_X86_64_APPLE_DARWIN_RUSTFLAGS" to "-C link-arg=-undefined -C link-arg=dynamic_lookup",
+            "CARGO_TARGET_AARCH64_APPLE_DARWIN_RUSTFLAGS" to "-C link-arg=-undefined -C link-arg=dynamic_lookup",
+        ),
+    )

codegen-server/python/src/test/kotlin/software/amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerTypesTest.kt

@@ -0,0 +1,147 @@
+/*
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ * SPDX-License-Identifier: Apache-2.0
+ */
+
+package software.amazon.smithy.rust.codegen.server.python.smithy.generators
+
+import org.junit.jupiter.api.Test
+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.testutil.asSmithyModel
+import software.amazon.smithy.rust.codegen.core.testutil.tokioTest
+import software.amazon.smithy.rust.codegen.core.util.dq
+import software.amazon.smithy.rust.codegen.server.python.smithy.testutil.cargoTest
+import software.amazon.smithy.rust.codegen.server.python.smithy.testutil.executePythonServerCodegenVisitor
+import software.amazon.smithy.rust.codegen.server.python.smithy.testutil.generatePythonServerPluginContext
+import kotlin.io.path.appendText
+
+internal class PythonServerTypesTest {
+    @Test
+    fun `document type`() {
+        val model = """
+            namespace test
+
+            use aws.protocols#restJson1
+
+            @restJson1
+            service Service {
+                operations: [
+                    Echo,
+                ],
+            }
+
+            @http(method: "POST", uri: "/echo")
+            operation Echo {
+                input: EchoInput,
+                output: EchoOutput,
+            }
+
+            structure EchoInput {
+                value: Document,
+            }
+
+            structure EchoOutput {
+                value: Document,
+            }
+        """.asSmithyModel()
+
+        val (pluginCtx, testDir) = generatePythonServerPluginContext(model)
+        executePythonServerCodegenVisitor(pluginCtx)
+
+        val testCases = listOf(
+            Pair(
+                """ { "value": 42 } """,
+                """
+                    assert input.value == 42
+                    output = EchoOutput(value=input.value)
+                """,
+            ),
+            Pair(
+                """ { "value": "foobar" } """,
+                """
+                    assert input.value == "foobar"
+                    output = EchoOutput(value=input.value)
+                """,
+            ),
+            Pair(
+                """
+                    {
+                        "value": [
+                            true,
+                            false,
+                            42,
+                            42.0,
+                            -42,
+                            {
+                                "nested": "value"
+                            },
+                            {
+                                "nested": [1, 2, 3]
+                            }
+                        ]
+                    }
+                """,
+                """
+                    assert input.value == [True, False, 42, 42.0, -42, {"nested": "value"}, {"nested": [1, 2, 3]}]
+                    output = EchoOutput(value=input.value)
+                """,
+            ),
+        )
+
+        val writer = RustWriter.forModule("service")
+        writer.tokioTest("document_type") {
+            rust(
+                """
+                use tower::Service as _;
+                use pyo3::{types::IntoPyDict, IntoPy, Python};
+                use hyper::{Body, Request, body};
+                use crate::{input, output};
+
+                pyo3::prepare_freethreaded_python();
+                """.trimIndent(),
+            )
+
+            testCases.forEach {
+                val payload = it.first.replace(" ", "").replace("\n", "")
+                val pythonHandler = it.second.trimIndent()
+                rust(
+                    """
+                    let mut service = Service::builder_without_plugins()
+                        .echo(|input: input::EchoInput| async {
+                            Ok(Python::with_gil(|py| {
+                                let globals = [("EchoOutput", py.get_type::<output::EchoOutput>())].into_py_dict(py);
+                                let locals = [("input", input.into_py(py))].into_py_dict(py);
+
+                                py.run(${pythonHandler.dq()}, Some(globals), Some(locals)).unwrap();
+
+                                locals
+                                    .get_item("output")
+                                    .unwrap()
+                                    .extract::<output::EchoOutput>()
+                                    .unwrap()
+                            }))
+                        })
+                        .build()
+                        .unwrap();
+
+                    let req = Request::builder()
+                        .method("POST")
+                        .uri("/echo")
+                        .body(Body::from(${payload.dq()}))
+                        .unwrap();
+
+                    let res = service.call(req).await.unwrap();
+                    assert!(res.status().is_success());
+                    let body = body::to_bytes(res.into_body()).await.unwrap();
+                    assert_eq!(body, ${payload.dq()});
+                    """.trimIndent(),
+                )
+            }
+        }
+
+        testDir.resolve("src/service.rs").appendText(writer.toString())
+
+        cargoTest(testDir)
+    }
+}