[Python] Improve Python stubs generation (#2606)

## Motivation and Context
This PR improves the Python stubs generation.

## Description
The main change is about avoiding to setup a placeholder for the Python
module and use the real module name, which allows to generate correct
docstrings during codegeneration.

We also change the stubs layout on disk, with the main stub entrypoint
called `__init__.pyi` instead of `$module_name.pyi`.

The README from the Rust runtime crate has been moved completely to the
example folder and I run autoformatting and style checks on the Python
example code.

----

_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 <1781140+crisidev@users.noreply.github.com>
Co-authored-by: Burak <unexge@gmail.com>

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

@@ -152,7 +152,7 @@ class PythonServerCodegenVisitor(
         rustCrate.useShapeWriter(shape) {
             // Use Python specific structure generator that adds the #[pyclass] attribute
             // and #[pymethods] implementation.
-            PythonServerStructureGenerator(model, codegenContext.symbolProvider, this, shape).render()
+            PythonServerStructureGenerator(model, codegenContext, this, shape).render()
 
             shape.getTrait<ErrorTrait>()?.also { errorTrait ->
                 ErrorImplGenerator(
@@ -190,7 +190,7 @@ class PythonServerCodegenVisitor(
     override fun unionShape(shape: UnionShape) {
         logger.info("[python-server-codegen] Generating an union shape $shape")
         rustCrate.useShapeWriter(shape) {
-            PythonServerUnionGenerator(model, codegenContext.symbolProvider, this, shape, renderUnknownVariant = false).render()
+            PythonServerUnionGenerator(model, codegenContext, this, shape, renderUnknownVariant = false).render()
         }
 
         if (shape.isReachableFromOperationInput() && shape.canReachConstrainedShape(

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

@@ -100,17 +100,14 @@ sealed class PythonType {
         override val namespace = type.namespace
     }
 
-    data class Opaque(override val name: String, val rustNamespace: String? = null) : PythonType() {
-        // Since Python doesn't have a something like Rust's `crate::` we are using a custom placeholder here
-        // and in our stub generation script we will replace placeholder with the real root module name.
-        private val pythonRootModulePlaceholder = "__root_module_name__"
+    data class Opaque(override val name: String, val pythonRootModuleName: String, val rustNamespace: String? = null) : PythonType() {
 
         override val namespace: String? = rustNamespace?.split("::")?.joinToString(".") {
             when (it) {
-                "crate" -> pythonRootModulePlaceholder
+                "crate" -> pythonRootModuleName
                 // In Python, we expose submodules from `aws_smithy_http_server_python`
-                // like `types`, `middleware`, `tls` etc. from `__root_module__name`
-                "aws_smithy_http_server_python" -> pythonRootModulePlaceholder
+                // like `types`, `middleware`, `tls` etc. from Python root module
+                "aws_smithy_http_server_python" -> pythonRootModuleName
                 else -> it
             }
         }
@@ -120,26 +117,29 @@ sealed class PythonType {
 /**
  * Return corresponding [PythonType] for a [RustType].
  */
-fun RustType.pythonType(): PythonType =
+fun RustType.pythonType(pythonRootModuleName: String): PythonType =
     when (this) {
         is RustType.Unit -> PythonType.None
         is RustType.Bool -> PythonType.Bool
         is RustType.Float -> PythonType.Float
         is RustType.Integer -> PythonType.Int
         is RustType.String -> PythonType.Str
-        is RustType.Vec -> PythonType.List(this.member.pythonType())
-        is RustType.Slice -> PythonType.List(this.member.pythonType())
-        is RustType.HashMap -> PythonType.Dict(this.key.pythonType(), this.member.pythonType())
-        is RustType.HashSet -> PythonType.Set(this.member.pythonType())
-        is RustType.Reference -> this.member.pythonType()
-        is RustType.Option -> PythonType.Optional(this.member.pythonType())
-        is RustType.Box -> this.member.pythonType()
-        is RustType.Dyn -> this.member.pythonType()
-        is RustType.Application -> PythonType.Application(this.type.pythonType(), this.args.map { it.pythonType() })
-        is RustType.Opaque -> PythonType.Opaque(this.name, this.namespace)
-        // TODO(Constraints): How to handle this?
-        // Revisit as part of https://github.com/awslabs/smithy-rs/issues/2114
-        is RustType.MaybeConstrained -> this.member.pythonType()
+        is RustType.Vec -> PythonType.List(this.member.pythonType(pythonRootModuleName))
+        is RustType.Slice -> PythonType.List(this.member.pythonType(pythonRootModuleName))
+        is RustType.HashMap -> PythonType.Dict(this.key.pythonType(pythonRootModuleName), this.member.pythonType(pythonRootModuleName))
+        is RustType.HashSet -> PythonType.Set(this.member.pythonType(pythonRootModuleName))
+        is RustType.Reference -> this.member.pythonType(pythonRootModuleName)
+        is RustType.Option -> PythonType.Optional(this.member.pythonType(pythonRootModuleName))
+        is RustType.Box -> this.member.pythonType(pythonRootModuleName)
+        is RustType.Dyn -> this.member.pythonType(pythonRootModuleName)
+        is RustType.Application -> PythonType.Application(
+            this.type.pythonType(pythonRootModuleName),
+            this.args.map {
+                it.pythonType(pythonRootModuleName)
+            },
+        )
+        is RustType.Opaque -> PythonType.Opaque(this.name, pythonRootModuleName, rustNamespace = this.namespace)
+        is RustType.MaybeConstrained -> this.member.pythonType(pythonRootModuleName)
     }
 
 /**

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

@@ -250,11 +250,11 @@ class PythonApplicationGenerator(
             """,
             *codegenScope,
         ) {
-            val middlewareRequest = PythonType.Opaque("Request", "crate::middleware")
-            val middlewareResponse = PythonType.Opaque("Response", "crate::middleware")
+            val middlewareRequest = PythonType.Opaque("Request", libName, rustNamespace = "crate::middleware")
+            val middlewareResponse = PythonType.Opaque("Response", libName, rustNamespace = "crate::middleware")
             val middlewareNext = PythonType.Callable(listOf(middlewareRequest), PythonType.Awaitable(middlewareResponse))
             val middlewareFunc = PythonType.Callable(listOf(middlewareRequest, middlewareNext), PythonType.Awaitable(middlewareResponse))
-            val tlsConfig = PythonType.Opaque("TlsConfig", "crate::tls")
+            val tlsConfig = PythonType.Opaque("TlsConfig", libName, rustNamespace = "crate::tls")
 
             rustTemplate(
                 """
@@ -344,9 +344,9 @@ class PythonApplicationGenerator(
                 val operationName = symbolProvider.toSymbol(operation).name
                 val fnName = RustReservedWords.escapeIfNeeded(symbolProvider.toSymbol(operation).name.toSnakeCase())
 
-                val input = PythonType.Opaque("${operationName}Input", "crate::input")
-                val output = PythonType.Opaque("${operationName}Output", "crate::output")
-                val context = PythonType.Opaque("Ctx")
+                val input = PythonType.Opaque("${operationName}Input", libName, rustNamespace = "crate::input")
+                val output = PythonType.Opaque("${operationName}Output", libName, rustNamespace = "crate::output")
+                val context = PythonType.Opaque("Ctx", libName)
                 val returnType = PythonType.Union(listOf(output, PythonType.Awaitable(output)))
                 val handler = PythonType.Union(
                     listOf(

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

@@ -18,16 +18,17 @@ import software.amazon.smithy.rust.codegen.core.rustlang.rust
 import software.amazon.smithy.rust.codegen.core.rustlang.rustInlineTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.rustTemplate
 import software.amazon.smithy.rust.codegen.core.rustlang.writable
-import software.amazon.smithy.rust.codegen.core.smithy.RustSymbolProvider
 import software.amazon.smithy.rust.codegen.core.smithy.generators.StructureGenerator
 import software.amazon.smithy.rust.codegen.core.smithy.rustType
 import software.amazon.smithy.rust.codegen.core.util.hasTrait
 import software.amazon.smithy.rust.codegen.core.util.isEventStream
+import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonEventStreamSymbolProvider
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerCargoDependency
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonType
 import software.amazon.smithy.rust.codegen.server.python.smithy.pythonType
 import software.amazon.smithy.rust.codegen.server.python.smithy.renderAsDocstring
+import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
 
 /**
  * To share structures defined in Rust with Python, `pyo3` provides the `PyClass` trait.
@@ -36,11 +37,13 @@ import software.amazon.smithy.rust.codegen.server.python.smithy.renderAsDocstrin
  */
 class PythonServerStructureGenerator(
     model: Model,
-    private val symbolProvider: RustSymbolProvider,
+    private val codegenContext: ServerCodegenContext,
     private val writer: RustWriter,
     private val shape: StructureShape,
-) : StructureGenerator(model, symbolProvider, writer, shape, emptyList()) {
+) : StructureGenerator(model, codegenContext.symbolProvider, writer, shape, emptyList()) {
 
+    private val symbolProvider = codegenContext.symbolProvider
+    private val libName = codegenContext.settings.moduleName.toSnakeCase()
     private val pyO3 = PythonServerCargoDependency.PyO3.toType()
 
     override fun renderStructure() {
@@ -157,9 +160,9 @@ class PythonServerStructureGenerator(
     private fun memberPythonType(shape: MemberShape, symbol: Symbol): PythonType =
         if (shape.isEventStream(model)) {
             val eventStreamSymbol = PythonEventStreamSymbolProvider.parseSymbol(symbol)
-            val innerT = eventStreamSymbol.innerT.pythonType()
+            val innerT = eventStreamSymbol.innerT.pythonType(libName)
             PythonType.AsyncIterator(innerT)
         } else {
-            symbol.rustType().pythonType()
+            symbol.rustType().pythonType(libName)
         }
 }

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

@@ -26,6 +26,7 @@ import software.amazon.smithy.rust.codegen.core.util.toSnakeCase
 import software.amazon.smithy.rust.codegen.server.python.smithy.PythonServerCargoDependency
 import software.amazon.smithy.rust.codegen.server.python.smithy.pythonType
 import software.amazon.smithy.rust.codegen.server.python.smithy.renderAsDocstring
+import software.amazon.smithy.rust.codegen.server.smithy.ServerCodegenContext
 
 /*
  * Generate unions that are compatible with Python by wrapping the Rust implementation into
@@ -34,11 +35,13 @@ import software.amazon.smithy.rust.codegen.server.python.smithy.renderAsDocstrin
  */
 class PythonServerUnionGenerator(
     model: Model,
-    private val symbolProvider: SymbolProvider,
+    private val codegenContext: ServerCodegenContext,
     private val writer: RustWriter,
     shape: UnionShape,
     private val renderUnknownVariant: Boolean = true,
-) : UnionGenerator(model, symbolProvider, writer, shape, renderUnknownVariant) {
+) : UnionGenerator(model, codegenContext.symbolProvider, writer, shape, renderUnknownVariant) {
+    private val symbolProvider = codegenContext.symbolProvider
+    private val libName = codegenContext.settings.moduleName.toSnakeCase()
     private val sortedMembers: List<MemberShape> = shape.allMembers.values.sortedBy { symbolProvider.toMemberName(it) }
     private val unionSymbol = symbolProvider.toSymbol(shape)
 
@@ -125,7 +128,7 @@ class PythonServerUnionGenerator(
             }
         } else {
             val memberSymbol = symbolProvider.toSymbol(member)
-            val pythonType = memberSymbol.rustType().pythonType()
+            val pythonType = memberSymbol.rustType().pythonType(libName)
             val targetType = memberSymbol.rustType()
             Attribute("staticmethod").render(writer)
             writer.rust(
@@ -166,7 +169,7 @@ class PythonServerUnionGenerator(
             }
         } else {
             val memberSymbol = symbolProvider.toSymbol(member)
-            val pythonType = memberSymbol.rustType().pythonType()
+            val pythonType = memberSymbol.rustType().pythonType(libName)
             val targetSymbol = symbolProvider.toSymbol(model.expectShape(member.target))
             val rustType = memberSymbol.rustType()
             writer.rust(
@@ -181,12 +184,13 @@ class PythonServerUnionGenerator(
                 } else {
                     "variant.clone()"
                 }
+                val errorVariant = memberSymbol.rustType().pythonType(libName).renderAsDocstring()
                 rustTemplate(
                     """
                     match self.0.as_$funcNamePart() {
                         Ok(variant) => Ok($variantType),
                         Err(_) => Err(#{pyo3}::exceptions::PyValueError::new_err(
-                            r"${unionSymbol.name} variant is not of type ${memberSymbol.rustType().pythonType().renderAsDocstring()}"
+                            r"${unionSymbol.name} variant is not of type $errorVariant"
                         )),
                     }
                     """,

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

@@ -28,9 +28,8 @@ internal class PythonTypeInformationGenerationTest {
         val foo = model.lookup<StructureShape>("test#Foo")
 
         val codegenContext = serverTestCodegenContext(model)
-        val symbolProvider = codegenContext.symbolProvider
         val writer = RustWriter.forModule("model")
-        PythonServerStructureGenerator(model, symbolProvider, writer, foo).render()
+        PythonServerStructureGenerator(model, codegenContext, writer, foo).render()
 
         val result = writer.toString()
 

examples/python/README.md

@@ -12,8 +12,8 @@ The Python implementation of the service can be found inside
     * [Python stub generation](#python-stub-generation)
 * [Run](#run)
 * [Test](#test)
-* [Uvloop](#uvloop)
 * [MacOs](#macos)
+* [Running servers on AWS Lambda](#running-servers-on-aws-lambda)
 
 ## Build
 
@@ -26,22 +26,28 @@ build.
 
 Ensure these dependencies are installed.
 
-```
+```bash
 pip install maturin uvloop aiohttp mypy
 ```
 
+The server can depend on [uvloop](https://pypi.org/project/uvloop/) for a faster
+event loop implementation and it will be automatically used instead of the standard
+library event loop if it is found in the dependencies' closure.
+
 ### 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
+* `make release`: build the Maturin package in release mode (includes type stubs
   generation).
-* `make install`: install the latest release or debug build.
+* `make install-wheel`: install the latest release or debug wheel using `pip`.
 * `make run`: run the example server.
 * `make test`: run the end-to-end integration test.
+* `make clean`: clean the Cargo artefacts.
+* `make dist-clean`: clean the Cargo artefacts and generated folders.
 
 ### Python stub generation
 
@@ -52,7 +58,7 @@ It will do first compilation of the crate, generate the types and exit.
 
 The script takes some command line arguments:
 
-```
+```bash
 ./stubgen.sh module_name manifest_path output_directory
 ```
 
@@ -72,23 +78,80 @@ 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.
+
+## Running servers on AWS Lambda
+
+`aws-smithy-http-server-python` supports running your services on [AWS Lambda](https://aws.amazon.com/lambda/).
+
+You need to use `run_lambda` method instead of `run` method to start
+the [custom runtime](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html)
+instead of the [Hyper](https://hyper.rs/) HTTP server.
+
+In your `app.py`:
+
+```diff
+from pokemon_service_server_sdk import App
+from pokemon_service_server_sdk.error import ResourceNotFoundException
+
+# ...
+
+# Get the number of requests served by this server.
+@app.get_server_statistics
+def get_server_statistics(
+    _: GetServerStatisticsInput, context: Context
+) -> GetServerStatisticsOutput:
+    calls_count = context.get_calls_count()
+    logging.debug("The service handled %d requests", calls_count)
+    return GetServerStatisticsOutput(calls_count=calls_count)
+
+# ...
+
+-app.run()
++app.run_lambda()
+```
+
+`aws-smithy-http-server-python` comes with a
+[custom runtime](https://docs.aws.amazon.com/lambda/latest/dg/runtimes-custom.html)
+so you should run your service without any provided runtimes.
+You can achieve that with a `Dockerfile` similar to this:
+
+```dockerfile
+# You can use any image that has your desired Python version
+FROM public.ecr.aws/lambda/python:3.8-x86_64
+
+# Copy your application code to `LAMBDA_TASK_ROOT`
+COPY app.py ${LAMBDA_TASK_ROOT}
+
+# When you build your Server SDK for your service, you will get a Python wheel.
+# You just need to copy that wheel and install it via `pip` inside your image.
+# Note that you need to build your library for Linux, and Python version used to
+# build your SDK should match with your image's Python version.
+# For cross compiling, you can consult to:
+# https://pyo3.rs/latest/building_and_distribution.html#cross-compiling
+COPY wheels/ ${LAMBDA_TASK_ROOT}/wheels
+RUN pip3 install ${LAMBDA_TASK_ROOT}/wheels/*.whl
+
+# You can install your application's other dependencies listed in `requirements.txt`.
+COPY requirements.txt .
+RUN pip3 install -r requirements.txt --target "${LAMBDA_TASK_ROOT}"
+
+# Create a symlink for your application's entrypoint,
+# so we can use `/app.py` to refer it
+RUN ln -s ${LAMBDA_TASK_ROOT}/app.py /app.py
+
+# By default `public.ecr.aws/lambda/python` images comes with Python runtime,
+# we need to override `ENTRYPOINT` and `CMD` to not call that runtime and
+# instead run directly your service and it will start our custom runtime.
+ENTRYPOINT [ "/var/lang/bin/python3.8" ]
+CMD [ "/app.py" ]
+```
+
+See [https://docs.aws.amazon.com/lambda/latest/dg/images-create.html#images-create-from-base](https://docs.aws.amazon.com/lambda/latest/dg/images-create.html#images-create-from-base)
+for more details on building your custom image.