[Python] Automatically generate stubs (#2576)

## Motivation and Context We want to automatically generate stubs in the codegen diff to ensure they can be reviewed and have a simple way to generate and include the stubs inside the Maturin wheel. ## Description The Python example has been moved to the `examples` folder and refactored. The refactoring ensures the script `stubgen.py` is included in the codegeneration of the SDK crate. The script is later used to generate stubs automatically during testing and can be used by customers to add their own stubs before the Maturin build ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ ---- _By submitting this pull request, I confirm that you can use, modify, copy, and redistribute this contribution, under the terms of your choice._ --------- Signed-off-by: Bigo <[email protected]> Co-authored-by: Burak <[email protected]>
smithy-lang · Apr 24, 2023 · 1a2dcd3 · 1a2dcd3
1 parent 81011a9
commit 1a2dcd3
Show file tree

Hide file tree

Showing 27 changed files with 291 additions and 139 deletions.
diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/rustlang/RustWriter.kt
@@ -469,6 +469,7 @@ class RustWriter private constructor(
                     devDependenciesOnly = true,
                 )
                 fileName == "package.json" -> rawWriter(fileName, debugMode = debugMode)
+                fileName == "stubgen.sh" -> rawWriter(fileName, debugMode = debugMode)
                 else -> RustWriter(fileName, namespace, debugMode = debugMode)
             }
         }

diff --git a/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt b/codegen-core/src/main/kotlin/software/amazon/smithy/rust/codegen/core/smithy/RuntimeType.kt
@@ -96,6 +96,8 @@ data class RuntimeConfig(
 
     val crateSrcPrefix: String = cratePrefix.replace("-", "_")
 
+    fun runtimeCratesPath(): String? = runtimeCrateLocation.path
+
     fun smithyRuntimeCrate(
         runtimeCrateName: String,
         optional: Boolean = false,

diff --git a/codegen-server-test/python/build.gradle.kts b/codegen-server-test/python/build.gradle.kts
@@ -66,7 +66,7 @@ val allCodegenTests = "../../codegen-core/common-test-models".let { commonModels
             "rest_json_extras",
             imports = listOf("$commonModels/rest-json-extras.smithy"),
         ),
-        // TODO(https://github.com/awslabs/smithy-rs/issues/2551)
+        // TODO(https://github.com/awslabs/smithy-rs/issues/2477)
         // CodegenTest(
         //     "aws.protocoltests.restjson.validation#RestJsonValidation",
         //     "rest_json_validation",
@@ -104,6 +104,21 @@ project.registerGenerateSmithyBuildTask(rootProject, pluginName, allCodegenTests
 project.registerGenerateCargoWorkspaceTask(rootProject, pluginName, allCodegenTests, workingDirUnderBuildDir)
 project.registerGenerateCargoConfigTomlTask(buildDir.resolve(workingDirUnderBuildDir))
 
+tasks.register("stubs") {
+    description = "Generate Python stubs for all models"
+    dependsOn("assemble")
+
+    doLast {
+        allCodegenTests.forEach { test ->
+            val crateDir = "$buildDir/$workingDirUnderBuildDir/${test.module}/$pluginName"
+            val moduleName = test.module.replace("-", "_")
+            exec {
+                commandLine("bash", "$crateDir/stubgen.sh", moduleName, "$crateDir/Cargo.toml", "$crateDir/python/$moduleName")
+            }
+        }
+    }
+}
+
 tasks["smithyBuildJar"].dependsOn("generateSmithyBuild")
 tasks["assemble"].finalizedBy("generateCargoWorkspace")
 

diff --git a/...n/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCargoDependency.kt b/...n/software/amazon/smithy/rust/codegen/server/python/smithy/PythonServerCargoDependency.kt
@@ -15,8 +15,8 @@ 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"))
-    val PyO3Asyncio: CargoDependency = CargoDependency("pyo3-asyncio", CratesIo("0.17"), features = setOf("attributes", "tokio-runtime"))
+    val PyO3: CargoDependency = CargoDependency("pyo3", CratesIo("0.18"))
+    val PyO3Asyncio: CargoDependency = CargoDependency("pyo3-asyncio", CratesIo("0.18"), features = setOf("attributes", "tokio-runtime"))
     val Tokio: CargoDependency = CargoDependency("tokio", CratesIo("1.20.1"), features = setOf("full"))
     val TokioStream: CargoDependency = CargoDependency("tokio-stream", CratesIo("0.1.12"))
     val Tracing: CargoDependency = CargoDependency("tracing", CratesIo("0.1"))

diff --git a/...n/smithy/rust/codegen/server/python/smithy/customizations/PythonServerCodegenDecorator.kt b/...n/smithy/rust/codegen/server/python/smithy/customizations/PythonServerCodegenDecorator.kt
@@ -23,6 +23,7 @@ import software.amazon.smithy.rust.codegen.server.python.smithy.generators.Pytho
 import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
 import software.amazon.smithy.rust.codegen.server.smithy.customizations.AddInternalServerErrorToAllOperationsDecorator
 import software.amazon.smithy.rust.codegen.server.smithy.customize.ServerCodegenDecorator
+import java.io.File
 
 /**
  * Configure the [lib] section of `Cargo.toml`.
@@ -194,6 +195,31 @@ class PyTypedMarkerDecorator : ServerCodegenDecorator {
     }
 }
 
+/**
+ * Copies the stubgen scripts to the generated crate root.
+ *
+ * The shell script `stubgen.sh` runs a quick build and uses `stubgen.py` to generate mypy compatibile
+ * types stubs for the project.
+ */
+class AddStubgenScriptDecorator : ServerCodegenDecorator {
+    override val name: String = "AddStubgenScriptDecorator"
+    override val order: Byte = 0
+
+    override fun extras(codegenContext: ServerCodegenContext, rustCrate: RustCrate) {
+        val runtimeCratesPath = codegenContext.runtimeConfig.runtimeCratesPath()
+        val stubgenPythonLocation = "$runtimeCratesPath/aws-smithy-http-server-python/stubgen.py"
+        val stubgenPythonContent = File(stubgenPythonLocation).readText(Charsets.UTF_8)
+        rustCrate.withFile("stubgen.py") {
+            writeWithNoFormatting("$stubgenPythonContent")
+        }
+        val stubgenShellLocation = "$runtimeCratesPath/aws-smithy-http-server-python/stubgen.sh"
+        val stubgenShellContent = File(stubgenShellLocation).readText(Charsets.UTF_8)
+        rustCrate.withFile("stubgen.sh") {
+            writeWithNoFormatting("$stubgenShellContent")
+        }
+    }
+}
+
 val DECORATORS = arrayOf(
     /**
      * Add the [InternalServerError] error to all operations.
@@ -214,4 +240,6 @@ val DECORATORS = arrayOf(
     InitPyDecorator(),
     // Generate `py.typed` for the Python source.
     PyTypedMarkerDecorator(),
+    // Generate scripts for stub generation.
+    AddStubgenScriptDecorator(),
 )
diff --git a/...amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerModuleGenerator.kt b/...amazon/smithy/rust/codegen/server/python/smithy/generators/PythonServerModuleGenerator.kt
@@ -11,7 +11,9 @@ import software.amazon.smithy.model.shapes.ResourceShape
 import software.amazon.smithy.model.shapes.ServiceShape
 import software.amazon.smithy.model.shapes.Shape
 import software.amazon.smithy.model.shapes.UnionShape
+import software.amazon.smithy.rust.codegen.core.Version
 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.rustlang.rustBlockTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.smithy.RustCrate
@@ -51,6 +53,8 @@ class PythonServerModuleGenerator(
                 renderPyTlsTypes()
                 renderPyLambdaTypes()
                 renderPyApplicationType()
+                renderCodegenVersion()
+                rust("Ok(())")
             }
         }
     }
@@ -62,13 +66,23 @@ class PythonServerModuleGenerator(
             let input = #{pyo3}::types::PyModule::new(py, "input")?;
             let output = #{pyo3}::types::PyModule::new(py, "output")?;
             let error = #{pyo3}::types::PyModule::new(py, "error")?;
-            let model = #{pyo3}::types::PyModule::new(py, "model")?;
             """,
             *codegenScope,
         )
+        // The `model` type section can be unused in models like `simple`, so we accommodate for it.
+        var visitedModelType = false
         serviceShapes.forEach { shape ->
             val moduleType = moduleType(shape)
             if (moduleType != null) {
+                if (moduleType == "model" && !visitedModelType) {
+                    rustTemplate(
+                        """
+                        let model = #{pyo3}::types::PyModule::new(py, "model")?;
+                        """,
+                        *codegenScope,
+                    )
+                    visitedModelType = true
+                }
                 when (shape) {
                     is UnionShape -> rustTemplate(
                         """
@@ -93,11 +107,18 @@ class PythonServerModuleGenerator(
             m.add_submodule(output)?;
             #{pyo3}::py_run!(py, error, "import sys; sys.modules['$libName.error'] = error");
             m.add_submodule(error)?;
-            #{pyo3}::py_run!(py, model, "import sys; sys.modules['$libName.model'] = model");
-            m.add_submodule(model)?;
             """,
             *codegenScope,
         )
+        if (visitedModelType) {
+            rustTemplate(
+                """
+                #{pyo3}::py_run!(py, model, "import sys; sys.modules['$libName.model'] = model");
+                m.add_submodule(model)?;
+                """,
+                *codegenScope,
+            )
+        }
     }
 
     // Render wrapper types that are substituted to the ones coming from `aws_smithy_types`.
@@ -211,13 +232,12 @@ class PythonServerModuleGenerator(
 
     // Render Python application type.
     private fun RustWriter.renderPyApplicationType() {
-        rustTemplate(
-            """
-            m.add_class::<crate::python_server_application::App>()?;
-            Ok(())
-            """,
-            *codegenScope,
-        )
+        rust("""m.add_class::<crate::python_server_application::App>()?;""")
+    }
+
+    // Render the codegeneration version as module attribute.
+    private fun RustWriter.renderCodegenVersion() {
+        rust("""m.add("CODEGEN_VERSION", "${Version.crateVersion()}")?;""")
     }
 
     // Convert to symbol and check the namespace to figure out where they should be imported from.

diff --git a/...software/amazon/smithy/rust/codegen/server/smithy/RustCrateInlineModuleComposingWriter.kt b/...software/amazon/smithy/rust/codegen/server/smithy/RustCrateInlineModuleComposingWriter.kt
@@ -335,7 +335,9 @@ class InnerModule(private val moduleDocProvider: ModuleDocProvider, debugMode: B
             inlineWriter
         } else {
             check(inlineModuleAndWriter.inlineModule == lookForModule) {
-                "The two inline modules have the same name but different attributes on them."
+                """The two inline modules have the same name but different attributes on them:
+                1) ${inlineModuleAndWriter.inlineModule}
+                2) $lookForModule"""
             }
 
             inlineModuleAndWriter.writer

diff --git a/...hy-http-server-python/examples/.gitignore → examples/python/.gitignore b/...hy-http-server-python/examples/.gitignore → examples/python/.gitignore
diff --git a/...hy-http-server-python/examples/Cargo.toml → examples/python/Cargo.toml b/...hy-http-server-python/examples/Cargo.toml → examples/python/Cargo.toml
diff --git a/...ithy-http-server-python/examples/Makefile → examples/python/Makefile b/...ithy-http-server-python/examples/Makefile → examples/python/Makefile
@@ -35,25 +35,23 @@ install-wheel:
 	find $(WHEELS) -type f -name '*.whl' | xargs python3 -m pip install --user --force-reinstall
 
 generate-stubs:
-	python3 $(CUR_DIR)/stubgen.py pokemon_service_server_sdk $(SERVER_SDK_DST)/python/pokemon_service_server_sdk
+	bash $(SERVER_SDK_DST)/stubgen.sh pokemon_service_server_sdk $(SERVER_SDK_DST)/Cargo.toml $(SERVER_SDK_DST)/python/pokemon_service_server_sdk
 
 build: codegen
+	$(MAKE) generate-stubs
 	$(MAKE) build-wheel
-	$(MAKE) install-wheel
+
+release: codegen
 	$(MAKE) generate-stubs
 	$(MAKE) build-wheel-release
-	$(MAKE) install-wheel
 
 run: build
 	python3 $(CUR_DIR)/pokemon_service.py
 
-py-check: build
+py-check: install-wheel
 	python3 -m mypy pokemon_service.py
 
-py-test:
-	python3 stubgen_test.py
-
-test: build py-check py-test
+test: build py-check
 	cargo test
 
 clippy: codegen

diff --git a/examples/python/README.md b/examples/python/README.md
@@ -0,0 +1,94 @@
+# Smithy Rust/Python Server SDK example
+
+This folder contains an example service called Pokémon Service used to showcase
+the service framework Python bindings capabilities and to run benchmarks.
+
+The Python implementation of the service can be found inside
+[pokemon_service.py](./pokemon_service.py).
+
+* [Build](#build)
+    * [Build dependencies](#build-dependencies)
+    * [Makefile](#makefile)
+    * [Python stub generation](#python-stub-generation)
+* [Run](#run)
+* [Test](#test)
+* [Uvloop](#uvloop)
+* [MacOs](#macos)
+
+## Build
+
+Since this example requires both the server and client SDK to be code-generated
+from their [model](/codegen-server-test/model/pokemon.smithy), a Makefile is
+provided to build and run the service. Just run `make build` to prepare the first
+build.
+
+### Build dependencies
+
+Ensure these dependencies are installed.
+
+```
+pip install maturin uvloop aiohttp mypy
+```
+
+### Makefile
+
+The build logic is drive by the Makefile:
+
+* `make codegen`: run the codegenerator.
+* `make build`: build the Maturin package in debug mode (includes type stubs
+  generation).
+* `make release`: build the Maturin package in release mode (includes type stub
+  generation).
+* `make install`: install the latest release or debug build.
+* `make run`: run the example server.
+* `make test`: run the end-to-end integration test.
+
+### Python stub generation
+
+We support the generation of mypy python stubs and every SDK crate ships with
+a script called `stubgen.sh`. **Note that the script is not called
+automatically as part of the build**. We suggest users to call it after code generation.
+It will do first compilation of the crate, generate the types and exit.
+
+The script takes some command line arguments:
+
+```
+./stubgen.sh module_name manifest_path output_directory
+```
+
+* module_name: name of the Python module to generate stubs for, IE `pokemon_service_server_sdk`.
+* manifest_path: path for the crate manifest used to build the types.
+* output_directory: directory where to generate the stub hierarchy. **This
+  directory should be a folder `python/$module_name` in the root of the Maturin package.**
+
+## Run
+
+`make run` can be used to start the Pokémon service on `http://localhost:13734`.
+
+## Test
+
+`make test` can be used to spawn the Python service and run some simple integration
+tests against it.
+
+More info can be found in the `tests` folder of `pokemon-service-test` package.
+
+## Uvloop
+
+The server can depend on [uvloop](https://pypi.org/project/uvloop/) for a
+faster event loop implementation. Uvloop can be installed with your favourite
+package manager or by using pip:
+
+```sh
+pip instal uvloop
+```
+
+and it will be automatically used instead of the standard library event loop if
+it is found in the dependencies' closure.
+
+## MacOs
+
+To compile and test on MacOs, please follow the official PyO3 guidelines on how
+to [configure your linker](https://pyo3.rs/latest/building_and_distribution.html?highlight=rustflags#macos).
+
+Please note that the `.cargo/config.toml` with linkers override can be local to
+your project.
diff --git a/...ithy-http-server-python/examples/mypy.ini → examples/python/mypy.ini b/...ithy-http-server-python/examples/mypy.ini → examples/python/mypy.ini
diff --git a/.../examples/pokemon-service-test/Cargo.toml → ...es/python/pokemon-service-test/Cargo.toml b/.../examples/pokemon-service-test/Cargo.toml → ...es/python/pokemon-service-test/Cargo.toml
@@ -9,15 +9,15 @@ description = "Run tests against the Python server implementation"
 [dev-dependencies]
 rand = "0.8"
 async-stream = "0.3"
-command-group = "1.0"
+command-group = "2.1.0"
 tokio = { version = "1.20.1", features = ["full"] }
-serial_test = "0.9.0"
+serial_test = "2.0.0"
 rustls-pemfile = "1.0.1"
-tokio-rustls = "0.23.4"
-hyper-rustls = { version = "0.23.0", features = ["http2"] }
+tokio-rustls = "0.24.0"
+hyper-rustls = { version = "0.24.0", features = ["http2"] }
 
 # Local paths
-aws-smithy-client  = { path = "../../../aws-smithy-client/", features = ["rustls"] }
-aws-smithy-http = { path = "../../../aws-smithy-http/" }
-aws-smithy-types  = { path = "../../../aws-smithy-types/" }
+aws-smithy-client  = { path = "../../../rust-runtime/aws-smithy-client/", features = ["rustls"] }
+aws-smithy-http = { path = "../../../rust-runtime/aws-smithy-http/" }
+aws-smithy-types  = { path = "../../../rust-runtime/aws-smithy-types/" }
 pokemon-service-client = { path = "../pokemon-service-client/" }
diff --git a/...les/pokemon-service-test/tests/helpers.rs → ...hon/pokemon-service-test/tests/helpers.rs b/...les/pokemon-service-test/tests/helpers.rs → ...hon/pokemon-service-test/tests/helpers.rs
diff --git a/...ice-test/tests/simple_integration_test.rs → ...ice-test/tests/simple_integration_test.rs b/...ice-test/tests/simple_integration_test.rs → ...ice-test/tests/simple_integration_test.rs
diff --git a/...service-test/tests/testdata/localhost.crt → ...service-test/tests/testdata/localhost.crt b/...service-test/tests/testdata/localhost.crt → ...service-test/tests/testdata/localhost.crt
diff --git a/...service-test/tests/testdata/localhost.key → ...service-test/tests/testdata/localhost.key b/...service-test/tests/testdata/localhost.key → ...service-test/tests/testdata/localhost.key
diff --git a/...server-python/examples/pokemon_service.py → examples/python/pokemon_service.py b/...server-python/examples/pokemon_service.py → examples/python/pokemon_service.py
diff --git a/rust-runtime/aws-smithy-http-server-python/Cargo.toml b/rust-runtime/aws-smithy-http-server-python/Cargo.toml
@@ -22,9 +22,9 @@ bytes = "1.2"
 futures = "0.3"
 http = "0.2"
 hyper = { version = "0.14.20", features = ["server", "http1", "http2", "tcp", "stream"] }
-tls-listener = { version = "0.5.1", features = ["rustls", "hyper-h2"] }
+tls-listener = { version = "0.7.0", features = ["rustls", "hyper-h2"] }
 rustls-pemfile = "1.0.1"
-tokio-rustls = "0.23.4"
+tokio-rustls = "0.24.0"
 lambda_http = { version = "0.7.1" }
 # There is a breaking change in `lambda_runtime` between `0.7.0` and `0.7.1`,
 # and `lambda_http` depends on `0.7` which by default resolves to `0.7.1` but in our CI
@@ -34,10 +34,10 @@ lambda_runtime = { version = "0.7.1" }
 num_cpus = "1.13.1"
 parking_lot = "0.12.1"
 pin-project-lite = "0.2"
-pyo3 = "0.17.0"
-pyo3-asyncio = { version = "0.17.0", features = ["tokio-runtime"] }
+pyo3 = "0.18.2"
+pyo3-asyncio = { version = "0.18.0", features = ["tokio-runtime"] }
 signal-hook = { version = "0.3.14", features = ["extended-siginfo"] }
-socket2 = { version = "0.4.4", features = ["all"] }
+socket2 = { version = "0.5.2", features = ["all"] }
 thiserror = "1.0.32"
 tokio = { version = "1.20.1", features = ["full"] }
 tokio-stream = "0.1"
@@ -51,9 +51,9 @@ pretty_assertions = "1"
 futures-util = { version = "0.3.16", default-features = false }
 tower-test = "0.4"
 tokio-test = "0.4"
-pyo3-asyncio = { version = "0.17.0", features = ["testing", "attributes", "tokio-runtime", "unstable-streams"] }
+pyo3-asyncio = { version = "0.18.0", features = ["testing", "attributes", "tokio-runtime", "unstable-streams"] }
 rcgen = "0.10.0"
-hyper-rustls = { version = "0.23.1", features = ["http2"] }
+hyper-rustls = { version = "0.24.0", features = ["http2"] }
 
 # PyO3 Asyncio tests cannot use Cargo's default testing harness because `asyncio`
 # wants to control the main thread. So we need to use testing harness provided by `pyo3_asyncio`